Ensure py.typed is present in all packages (BUILD-660)

dagster-io · Feb 2, 2025 · 65ff6fc · 65ff6fc
1 parent 357f918
commit 65ff6fc
Show file tree

Hide file tree

Showing 60 changed files with 95 additions and 1 deletion.
diff --git a/.buildkite/dagster-buildkite/dagster_buildkite/py.typed b/.buildkite/dagster-buildkite/dagster_buildkite/py.typed
diff --git a/docs/dagster-ui-screenshot/dagster_ui_screenshot/py.typed b/docs/dagster-ui-screenshot/dagster_ui_screenshot/py.typed
diff --git a/docs/sphinx/_ext/dagster-sphinx/dagster_sphinx/py.typed b/docs/sphinx/_ext/dagster-sphinx/dagster_sphinx/py.typed
diff --git a/examples/airlift-federation-tutorial/airlift_federation_tutorial/py.typed b/examples/airlift-federation-tutorial/airlift_federation_tutorial/py.typed
diff --git a/examples/airlift-migration-tutorial/tutorial_example/py.typed b/examples/airlift-migration-tutorial/tutorial_example/py.typed
diff --git a/examples/airlift-mwaa-example/mwaa_example/py.typed b/examples/airlift-mwaa-example/mwaa_example/py.typed
diff --git a/examples/assets_dbt_python/assets_dbt_python/py.typed b/examples/assets_dbt_python/assets_dbt_python/py.typed
diff --git a/examples/assets_dynamic_partitions/assets_dynamic_partitions/py.typed b/examples/assets_dynamic_partitions/assets_dynamic_partitions/py.typed
diff --git a/examples/assets_modern_data_stack/assets_modern_data_stack/py.typed b/examples/assets_modern_data_stack/assets_modern_data_stack/py.typed
diff --git a/examples/assets_pandas_pyspark/assets_pandas_pyspark/py.typed b/examples/assets_pandas_pyspark/assets_pandas_pyspark/py.typed
diff --git a/examples/assets_pandas_type_metadata/assets_pandas_type_metadata/py.typed b/examples/assets_pandas_type_metadata/assets_pandas_type_metadata/py.typed
diff --git a/examples/assets_smoke_test/assets_smoke_test/py.typed b/examples/assets_smoke_test/assets_smoke_test/py.typed
diff --git a/examples/deploy_k8s_beta/py.typed b/examples/deploy_k8s_beta/py.typed
diff --git a/examples/development_to_production/development_to_production/py.typed b/examples/development_to_production/development_to_production/py.typed
diff --git a/examples/docs_beta_snippets/docs_beta_snippets/py.typed b/examples/docs_beta_snippets/docs_beta_snippets/py.typed
diff --git a/examples/docs_snippets/docs_snippets/py.typed b/examples/docs_snippets/docs_snippets/py.typed
diff --git a/examples/experimental/assets_yaml_dsl/assets_yaml_dsl/py.typed b/examples/experimental/assets_yaml_dsl/assets_yaml_dsl/py.typed
diff --git a/examples/experimental/dagster-airlift/examples/dbt-example/py.typed b/examples/experimental/dagster-airlift/examples/dbt-example/py.typed
diff --git a/examples/experimental/dagster-airlift/examples/kitchen-sink/py.typed b/examples/experimental/dagster-airlift/examples/kitchen-sink/py.typed
diff --git a/examples/experimental/dagster-blueprints/dagster_blueprints/py.typed b/examples/experimental/dagster-blueprints/dagster_blueprints/py.typed
diff --git a/...s/experimental/dagster-blueprints/examples/builtin-blueprints/builtin_blueprints/py.typed b/...s/experimental/dagster-blueprints/examples/builtin-blueprints/builtin_blueprints/py.typed
diff --git a/...les/experimental/dagster-blueprints/examples/custom-blueprints/custom_blueprints/py.typed b/...les/experimental/dagster-blueprints/examples/custom-blueprints/custom_blueprints/py.typed
diff --git a/examples/experimental/dagster-dlift/dagster_dlift/py.typed b/examples/experimental/dagster-dlift/dagster_dlift/py.typed
diff --git a/examples/experimental/dagster-dlift/kitchen-sink/dlift_kitchen_sink/py.typed b/examples/experimental/dagster-dlift/kitchen-sink/dlift_kitchen_sink/py.typed
diff --git a/examples/experimental/external_assets/external_assets/py.typed b/examples/experimental/external_assets/external_assets/py.typed
diff --git a/examples/feature_graph_backed_assets/feature_graph_backed_assets/py.typed b/examples/feature_graph_backed_assets/feature_graph_backed_assets/py.typed
diff --git a/examples/getting_started_etl_tutorial/py.typed b/examples/getting_started_etl_tutorial/py.typed
diff --git a/examples/project_analytics/dagster_pypi/py.typed b/examples/project_analytics/dagster_pypi/py.typed
diff --git a/examples/project_dagster_modal_pipes/project_dagster_modal_pipes/py.typed b/examples/project_dagster_modal_pipes/project_dagster_modal_pipes/py.typed
diff --git a/examples/project_dagster_university_start/dagster_university/py.typed b/examples/project_dagster_university_start/dagster_university/py.typed
diff --git a/examples/project_du_dbt_starter/dagster_university/py.typed b/examples/project_du_dbt_starter/dagster_university/py.typed
diff --git a/examples/project_fully_featured/project_fully_featured/py.typed b/examples/project_fully_featured/project_fully_featured/py.typed
diff --git a/examples/project_llm_fine_tune/project_llm_fine_tune/py.typed b/examples/project_llm_fine_tune/project_llm_fine_tune/py.typed
diff --git a/examples/project_prompt_eng/project_prompt_eng/py.typed b/examples/project_prompt_eng/project_prompt_eng/py.typed
diff --git a/examples/quickstart_aws/quickstart_aws/py.typed b/examples/quickstart_aws/quickstart_aws/py.typed
diff --git a/examples/quickstart_etl/quickstart_etl/py.typed b/examples/quickstart_etl/quickstart_etl/py.typed
diff --git a/examples/quickstart_gcp/quickstart_gcp/py.typed b/examples/quickstart_gcp/quickstart_gcp/py.typed
diff --git a/examples/quickstart_snowflake/quickstart_snowflake/py.typed b/examples/quickstart_snowflake/quickstart_snowflake/py.typed
diff --git a/examples/starlift-demo/dbt_example/py.typed b/examples/starlift-demo/dbt_example/py.typed
diff --git a/examples/tutorial/tutorial/py.typed b/examples/tutorial/tutorial/py.typed
diff --git a/examples/tutorial_notebook_assets/tutorial_finished/py.typed b/examples/tutorial_notebook_assets/tutorial_finished/py.typed
diff --git a/examples/use_case_repository/use_case_repository/py.typed b/examples/use_case_repository/use_case_repository/py.typed
diff --git a/examples/with_great_expectations/with_great_expectations/py.typed b/examples/with_great_expectations/with_great_expectations/py.typed
diff --git a/examples/with_openai/with_openai/py.typed b/examples/with_openai/with_openai/py.typed
diff --git a/examples/with_pyspark/with_pyspark/py.typed b/examples/with_pyspark/with_pyspark/py.typed
diff --git a/examples/with_pyspark_emr/with_pyspark_emr/py.typed b/examples/with_pyspark_emr/with_pyspark_emr/py.typed
diff --git a/examples/with_wandb/with_wandb/py.typed b/examples/with_wandb/with_wandb/py.typed
diff --git a/helm/dagster/schema/schema/py.typed b/helm/dagster/schema/schema/py.typed
diff --git a/integration_tests/python_modules/dagster-k8s-test-infra/dagster_k8s_test_infra/py.typed b/integration_tests/python_modules/dagster-k8s-test-infra/dagster_k8s_test_infra/py.typed
diff --git a/integration_tests/test_suites/dagster-azure-live-tests/azure_test_proj/py.typed b/integration_tests/test_suites/dagster-azure-live-tests/azure_test_proj/py.typed
diff --git a/integration_tests/test_suites/dagster-gcp-live-tests/gcp_test_proj/py.typed b/integration_tests/test_suites/dagster-gcp-live-tests/gcp_test_proj/py.typed
diff --git a/python_modules/automation/automation/utils.py b/python_modules/automation/automation/utils.py
@@ -1,6 +1,9 @@
+import os
 import subprocess
-from collections.abc import Iterable
+from collections.abc import Iterable, Iterator
+from contextlib import contextmanager
 from distutils import spawn
+from pathlib import Path
 from typing import Optional
 
 import click
@@ -24,3 +27,40 @@ def which_(exe: str) -> Optional[str]:
 
 def all_equal(iterable: Iterable[object]) -> bool:
     return len(set(iterable)) == 1
+
+
+def discover_git_root(path: Path) -> Path:
+    while path != path.parent:
+        if (path / ".git").exists():
+            return path
+        path = path.parent
+    raise ValueError("Could not find git root")
+
+
+@contextmanager
+def pushd(path: Path) -> Iterator[None]:
+    original_dir = Path.cwd()
+    os.chdir(path)
+    yield
+    os.chdir(original_dir)
+
+
+def git_ls_files(pattern: str) -> list[str]:
+    return (
+        subprocess.run(["git", "ls-files", pattern], check=True, text=True, capture_output=True)
+        .stdout.strip()
+        .split("\n")
+    )
+
+
+def get_all_repo_packages() -> list[Path]:
+    git_root = discover_git_root(Path(__file__))
+    with pushd(git_root):
+        return [
+            Path(p).parent
+            for p in subprocess.run(
+                ["git", "ls-files", "**/setup.py"], check=True, text=True, capture_output=True
+            )
+            .stdout.strip()
+            .split("\n")
+        ]
diff --git a/python_modules/automation/automation_tests/test_repo.py b/python_modules/automation/automation_tests/test_repo.py
@@ -2,6 +2,8 @@
 import subprocess
 from pathlib import Path
 
+from automation.utils import discover_git_root, get_all_repo_packages, git_ls_files, pushd
+
 # Some libraries are excluded because they either:
 # - lack a Dagster dependency, which is a prerequisite for registering in the DagsterLibraryRegistry.
 # - are temporary or on a separate release schedule from the rest of the libraries.
@@ -26,3 +28,54 @@ def test_all_libraries_register() -> None:
         assert (
             result.returncode == 0
         ), f"Dagster library {library} is missing call to {register_call}."
+
+
+# We want to make sure all of our packages are published with `py.typed` files unless explicitly
+# excluded.
+def test_all_packages_have_py_typed():
+    git_root = discover_git_root(Path(__file__))
+    with pushd(git_root):
+        package_roots = get_all_repo_packages()
+        missing_py_typed_file = []
+        missing_py_typed_in_manifest_in = []
+        for package_root in package_roots:
+            root_python_package = _find_root_python_package(package_root)
+            if not (root_python_package / "py.typed").exists():
+                missing_py_typed_file.append(str(package_root))
+
+            not_published_packages = [
+                "automation",
+                "dagster-test",
+                "kitchen-sink",  # in dagster-airlift
+                "perf-harness",  # in dagster-airlift
+            ]
+
+            # Published packages are additionally required to include py.typed in MANIFEST.in to ensure
+            # they are included in the published distribution.
+            if (
+                str(package_root).startswith("python_modules")
+                and package_root.name not in not_published_packages
+            ):
+                manifest_in = package_root / "MANIFEST.in"
+                assert manifest_in.exists(), f"MANIFEST.in is missing for: {package_root}"
+                if "py.typed" not in manifest_in.read_text():
+                    with open(manifest_in) as f:
+                        if "py.typed" not in f.read():
+                            missing_py_typed_in_manifest_in.append(str(package_root))
+
+        nl = "\n"
+        assert (
+            not missing_py_typed_file
+        ), f"Missing py.typed files in these packages:\n{nl.join(missing_py_typed_file)}"
+        assert not missing_py_typed_in_manifest_in, f"Missing py.typed in MANIFEST.in for these packages:\n{nl.join(missing_py_typed_in_manifest_in)}"
+
+
+def _find_root_python_package(package_root: Path) -> Path:
+    standard_name = package_root / package_root.name.replace("-", "_")
+    if standard_name.exists():
+        return standard_name
+    else:
+        with pushd(package_root):
+            init_pys = git_ls_files("*/__init__.py")
+            packages = [Path(p).parent for p in init_pys]
+            return next(package_root / p for p in packages if not p.name.endswith("_tests"))
diff --git a/python_modules/dagit/dagit/py.typed b/python_modules/dagit/dagit/py.typed
diff --git a/python_modules/dagster-test/dagster_test/py.typed b/python_modules/dagster-test/dagster_test/py.typed
diff --git a/python_modules/libraries/dagster-airlift/kitchen-sink/kitchen_sink/py.typed b/python_modules/libraries/dagster-airlift/kitchen-sink/kitchen_sink/py.typed
diff --git a/python_modules/libraries/dagster-airlift/perf-harness/perf_harness/py.typed b/python_modules/libraries/dagster-airlift/perf-harness/perf_harness/py.typed
diff --git a/python_modules/libraries/dagster-census/MANIFEST.in b/python_modules/libraries/dagster-census/MANIFEST.in
@@ -1,2 +1,3 @@
 include README.md
 include LICENSE
+include dagster_census/py.typed
diff --git a/python_modules/libraries/dagster-census/dagster_census/py.typed b/python_modules/libraries/dagster-census/dagster_census/py.typed
diff --git a/python_modules/libraries/dagster-managed-elements/dagster_managed_elements/py.typed b/python_modules/libraries/dagster-managed-elements/dagster_managed_elements/py.typed