[dg] Refactor PackageEntry types

Author: OwenKephart

python_modules/dagster/dagster/components/cli/list.py

@@ -26,7 +26,7 @@
     discover_entry_point_package_entries,
     discover_package_entries,
 )
-from dagster.components.core.snapshot import get_library_object_snap
+from dagster.components.core.snapshot import get_package_entry_snap
 
 
 @click.group(name="list")
@@ -44,7 +44,7 @@ def list_library_command(
     """List registered library objects."""
     library_objects = _load_library_objects(entry_points, extra_modules)
     serialized_snaps = serialize_value(
-        [get_library_object_snap(key, obj) for key, obj in library_objects.items()]
+        [get_package_entry_snap(key, obj) for key, obj in library_objects.items()]
     )
     click.echo(serialized_snaps)
 

python_modules/dagster/dagster/components/core/snapshot.py

@@ -3,8 +3,8 @@
 
 from dagster_shared.serdes.objects.package_entry import (
     ComponentTypeData,
-    LibraryEntryKey,
-    LibraryEntrySnap,
+    PackageEntryKey,
+    PackageEntrySnap,
     ScaffoldTargetTypeData,
 )
 
@@ -43,12 +43,12 @@ def _get_scaffold_target_type_data(scaffolder: Scaffolder) -> ScaffoldTargetType
     )
 
 
-def get_library_object_snap(key: LibraryEntryKey, obj: object) -> LibraryEntrySnap:
+def get_package_entry_snap(key: PackageEntryKey, obj: object) -> PackageEntrySnap:
     type_data = []
     if isinstance(obj, type) and issubclass(obj, Component):
         type_data.append(_get_component_type_data(obj))
     scaffolder = get_scaffolder(obj) if isinstance(obj, type) else None
     if isinstance(scaffolder, Scaffolder):
         type_data.append(_get_scaffold_target_type_data(scaffolder))
     summary, description = _get_summary_and_description(obj)
-    return LibraryEntrySnap(key=key, summary=summary, description=description, type_data=type_data)
+    return PackageEntrySnap(key=key, summary=summary, description=description, type_data=type_data)

python_modules/dagster/dagster_tests/components_tests/cli_tests/test_commands.py

@@ -7,10 +7,10 @@
 from dagster.components.cli import cli
 from dagster_shared import check
 from dagster_shared.serdes.objects import (
-    ComponentTypeSnap,
+    ComponentTypeData,
     PackageEntryKey,
     PackageEntrySnap,
-    ScaffolderSnap,
+    ScaffoldTargetTypeData,
 )
 from dagster_shared.serdes.serdes import deserialize_value
 from jsonschema import Draft202012Validator, ValidationError
@@ -50,21 +50,25 @@ def test_list_library_objects_from_module():
         "dagster_test.components.SimplePipesScriptComponent",
     ]
 
-    assert result[2] == ComponentTypeSnap(
+    assert result[2] == PackageEntrySnap(
         key=PackageEntryKey(namespace="dagster_test.components", name="SimpleAssetComponent"),
-        schema={
-            "additionalProperties": False,
-            "properties": {
-                "asset_key": {"title": "Asset Key", "type": "string"},
-                "value": {"title": "Value", "type": "string"},
-            },
-            "required": ["asset_key", "value"],
-            "title": "SimpleAssetComponentModel",
-            "type": "object",
-        },
         description="A simple asset that returns a constant string value.",
         summary="A simple asset that returns a constant string value.",
-        scaffolder=ScaffolderSnap(schema=None),
+        type_data=[
+            ComponentTypeData(
+                schema={
+                    "additionalProperties": False,
+                    "properties": {
+                        "asset_key": {"title": "Asset Key", "type": "string"},
+                        "value": {"title": "Value", "type": "string"},
+                    },
+                    "required": ["asset_key", "value"],
+                    "title": "SimpleAssetComponentModel",
+                    "type": "object",
+                }
+            ),
+            ScaffoldTargetTypeData(schema=None),
+        ],
     )
 
     pipes_script_params_schema = {
@@ -77,12 +81,14 @@ def test_list_library_objects_from_module():
         "type": "object",
     }
 
-    assert result[3] == ComponentTypeSnap(
+    assert result[3] == PackageEntrySnap(
         key=PackageEntryKey(namespace="dagster_test.components", name="SimplePipesScriptComponent"),
-        schema=pipes_script_params_schema,
         description="A simple asset that runs a Python script with the Pipes subprocess client.\n\nBecause it is a pipes asset, no value is returned.",
         summary="A simple asset that runs a Python script with the Pipes subprocess client.",
-        scaffolder=ScaffolderSnap(schema=pipes_script_params_schema),
+        type_data=[
+            ComponentTypeData(schema=pipes_script_params_schema),
+            ScaffoldTargetTypeData(schema=pipes_script_params_schema),
+        ],
     )
 
 

python_modules/dagster/dagster_tests/components_tests/registry_tests/test_registry.py

@@ -11,7 +11,7 @@
 from dagster._utils import pushd
 from dagster.components import Component
 from dagster.components.core.package_entry import discover_entry_point_package_entries
-from dagster.components.core.snapshot import get_library_object_snap
+from dagster.components.core.snapshot import get_package_entry_snap
 from dagster_dg.utils import get_venv_executable
 from dagster_shared.serdes.objects import PackageEntryKey
 from dagster_shared.serdes.serdes import deserialize_value
@@ -129,7 +129,7 @@ def test_all_components_have_defined_summary():
     registry = discover_entry_point_package_entries()
     for component_name, component_type in registry.items():
         if isinstance(component_type, type) and issubclass(component_type, Component):
-            assert get_library_object_snap(
+            assert get_package_entry_snap(
                 PackageEntryKey("a", "a"), component_type
             ).summary, f"Component {component_name} has no summary defined"
 

python_modules/libraries/dagster-dg/dagster_dg/check.py

@@ -117,7 +117,7 @@ def check_yaml(
     )
     for key, component_doc_tree in component_contents_by_key.items():
         try:
-            json_schema = component_registry.get_component_type(key).schema or {}
+            json_schema = component_registry.get(key).component_schema or {}
 
             v = Draft202012Validator(json_schema)
             for err in v.iter_errors(component_doc_tree.value["attributes"]):

python_modules/libraries/dagster-dg/dagster_dg/cli/list.py

@@ -1,10 +1,9 @@
 import json
 from dataclasses import asdict
 from pathlib import Path
-from typing import Any
+from typing import Any, Optional
 
 import click
-from dagster_shared.serdes.objects import ComponentTypeSnap
 from rich.console import Console
 from rich.table import Table
 
@@ -68,27 +67,17 @@ def list_component_command(**global_options: object) -> None:
 # ########################
 
 
-@list_group.command(name="component-type", cls=DgClickCommand)
-@click.option(
-    "--json",
-    "output_json",
-    is_flag=True,
-    default=False,
-    help="Output as JSON instead of a table.",
-)
-@dg_global_options
-@cli_telemetry_wrapper
-def list_component_type_command(output_json: bool, **global_options: object) -> None:
-    """List registered Dagster components in the current project environment."""
-    cli_config = normalize_cli_config(global_options, click.get_current_context())
-    dg_context = DgContext.for_defined_registry_environment(Path.cwd(), cli_config)
-    registry = RemotePackageRegistry.from_dg_context(dg_context)
-
+def _list_package_entries(
+    registry: RemotePackageRegistry, output_json: bool, entry_type: Optional[str]
+) -> None:
     sorted_keys = sorted(
-        (k for k in registry.keys() if isinstance(registry.get(k), ComponentTypeSnap)),
+        [
+            key
+            for key in registry.keys()
+            if entry_type is None or entry_type in registry.get(key).types
+        ],
         key=lambda k: k.to_typename(),
     )
-
     # JSON
     if output_json:
         output: list[dict[str, object]] = []
@@ -101,18 +90,36 @@ def list_component_type_command(output_json: bool, **global_options: object) ->
                 }
             )
         click.echo(json.dumps(output, indent=4))
-
     # TABLE
     else:
         table = Table(border_style="dim")
-        table.add_column("Component Type", style="bold cyan", no_wrap=True)
+        table.add_column("Package Entry", style="bold cyan", no_wrap=True)
         table.add_column("Summary")
         for key in sorted_keys:
             table.add_row(key.to_typename(), registry.get(key).summary)
         console = Console()
         console.print(table)
 
 
+@list_group.command(name="component-type", cls=DgClickCommand)
+@click.option(
+    "--json",
+    "output_json",
+    is_flag=True,
+    default=False,
+    help="Output as JSON instead of a table.",
+)
+@dg_global_options
+@cli_telemetry_wrapper
+def list_component_type_command(output_json: bool, **global_options: object) -> None:
+    """List registered Dagster components in the current project environment."""
+    cli_config = normalize_cli_config(global_options, click.get_current_context())
+    dg_context = DgContext.for_defined_registry_environment(Path.cwd(), cli_config)
+    registry = RemotePackageRegistry.from_dg_context(dg_context)
+
+    _list_package_entries(registry, output_json, "component")
+
+
 # ########################
 # ##### DEFS
 # ########################

python_modules/libraries/dagster-dg/dagster_dg/cli/utils.py

@@ -91,36 +91,34 @@ def inspect_component_type_command(
             "Only one of --description, --scaffold-params-schema, and --component-schema can be specified."
         )
 
-    component_type_snap = registry.get_component_type(component_key)
-
+    entry_snap = registry.get(component_key)
     if description:
-        if component_type_snap.description:
-            click.echo(component_type_snap.description)
+        if entry_snap.description:
+            click.echo(entry_snap.description)
         else:
             click.echo("No description available.")
     elif scaffold_params_schema:
-        if component_type_snap.scaffolder_schema:
-            click.echo(_serialize_json_schema(component_type_snap.scaffolder_schema))
+        if entry_snap.scaffolder_schema:
+            click.echo(_serialize_json_schema(entry_snap.scaffolder_schema))
         else:
             click.echo("No scaffold params schema defined.")
     elif component_schema:
-        if component_type_snap.schema:
-            click.echo(_serialize_json_schema(component_type_snap.schema))
+        if entry_snap.component_schema:
+            click.echo(_serialize_json_schema(entry_snap.component_schema))
         else:
             click.echo("No component schema defined.")
-
     # print all available metadata
     else:
         click.echo(component_type)
-        if component_type_snap.description:
+        if entry_snap.description:
             click.echo("\nDescription:\n")
-            click.echo(component_type_snap.description)
-        if component_type_snap.scaffolder_schema:
+            click.echo(entry_snap.description)
+        if entry_snap.scaffolder_schema:
             click.echo("\nScaffold params schema:\n")
-            click.echo(_serialize_json_schema(component_type_snap.scaffolder_schema))
-        if component_type_snap.schema:
+            click.echo(_serialize_json_schema(entry_snap.scaffolder_schema))
+        if entry_snap.component_schema:
             click.echo("\nComponent schema:\n")
-            click.echo(_serialize_json_schema(component_type_snap.schema))
+            click.echo(_serialize_json_schema(entry_snap.component_schema))
 
 
 def _serialize_json_schema(schema: Mapping[str, Any]) -> str:

python_modules/libraries/dagster-dg/dagster_dg/component.py

@@ -5,7 +5,7 @@
 
 import dagster_shared.check as check
 from dagster_shared.serdes import deserialize_value, serialize_value
-from dagster_shared.serdes.objects import ComponentTypeSnap, PackageEntryKey, PackageEntrySnap
+from dagster_shared.serdes.objects import PackageEntryKey, PackageEntrySnap
 
 from dagster_dg.utils import is_valid_json
 
@@ -48,13 +48,6 @@ def get(self, key: PackageEntryKey) -> PackageEntrySnap:
         """Resolves a library object within the scope of a given component directory."""
         return self._objects[key]
 
-    def get_component_type(self, key: PackageEntryKey) -> ComponentTypeSnap:
-        """Resolves a component type within the scope of a given component directory."""
-        obj = self.get(key)
-        if not isinstance(obj, ComponentTypeSnap):
-            raise ValueError(f"Expected component type, got {obj}")
-        return obj
-
     def has(self, key: PackageEntryKey) -> bool:
         return key in self._objects
 

python_modules/libraries/dagster-dg/dagster_dg/docs.py

@@ -5,11 +5,12 @@
 from typing import Any, Optional, TypedDict, Union
 
 import yaml
-from dagster_shared.serdes.objects import PackageEntryKey
+from dagster_shared.serdes.objects import ComponentTypeData, PackageEntryKey
+from dagster_shared.serdes.objects.package_entry import PackageEntrySnap
 from dagster_shared.yaml_utils import parse_yaml_with_source_positions
 from dagster_shared.yaml_utils.source_position import SourcePositionTree
 
-from dagster_dg.component import ComponentTypeSnap, RemotePackageRegistry
+from dagster_dg.component import RemotePackageRegistry
 
 REF_BASE = "#/$defs/"
 JSON_SCHEMA_EXTRA_REQUIRED_SCOPE_KEY = "dagster_required_scope"
@@ -175,11 +176,16 @@ def json_for_all_components(
     registry: RemotePackageRegistry,
 ) -> list[ComponentTypeNamespaceJson]:
     """Returns a list of JSON representations of all component types in the registry."""
-    component_json = [
-        (key.namespace.split(".")[0], json_for_component_type(key, library_obj))
-        for key, library_obj in registry.items()
-        if isinstance(library_obj, ComponentTypeSnap) and library_obj.schema is not None
-    ]
+    component_json = []
+    for key, entry in registry.items():
+        component_type_data = entry.get_type_data("component")
+        if component_type_data and component_type_data.schema:
+            component_json.append(
+                (
+                    key.namespace.split(".")[0],
+                    json_for_component_type(key, entry, component_type_data),
+                )
+            )
     return [
         ComponentTypeNamespaceJson(
             name=namespace,
@@ -190,15 +196,15 @@ def json_for_all_components(
 
 
 def json_for_component_type(
-    key: PackageEntryKey, remote_component_type: ComponentTypeSnap
+    key: PackageEntryKey, entry: PackageEntrySnap, component_type_data: ComponentTypeData
 ) -> ComponentTypeJson:
     typename = key.to_typename()
-    sample_yaml = generate_sample_yaml(typename, remote_component_type.schema or {})
+    sample_yaml = generate_sample_yaml(typename, component_type_data.schema or {})
     return ComponentTypeJson(
         name=typename,
         author="",
         tags=[],
         example=sample_yaml,
-        schema=json.dumps(remote_component_type.schema),
-        description=remote_component_type.description,
+        schema=json.dumps(component_type_data.schema),
+        description=entry.description,
     )

python_modules/libraries/dagster-shared/dagster_shared/serdes/objects/__init__.py

@@ -1,6 +1,6 @@
 from dagster_shared.serdes.objects.package_entry import (
-    ComponentTypeSnap as ComponentTypeSnap,
+    ComponentTypeData as ComponentTypeData,
     PackageEntryKey as PackageEntryKey,
     PackageEntrySnap as PackageEntrySnap,
-    ScaffolderSnap as ScaffolderSnap,
+    ScaffoldTargetTypeData as ScaffoldTargetTypeData,
 )