Skip to content

Commit 16e0a19

Browse files
OwenKephartOwenKephart
committed
[dg] dg list packages
1 parent 1a156b6 commit 16e0a19

File tree

3 files changed

+114
-34
lines changed

3 files changed

+114
-34
lines changed

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

Lines changed: 91 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,17 @@
11
import json
2+
from collections.abc import Sequence
23
from dataclasses import asdict
34
from pathlib import Path
45
from typing import Any, Optional
56

67
import click
8+
from dagster_shared.serdes.objects import PackageEntrySnap
79
from rich.console import Console
810
from rich.table import Table
11+
from rich.text import Text
912

1013
from dagster_dg.cli.shared_options import dg_global_options
11-
from dagster_dg.component import RemotePackageRegistry
14+
from dagster_dg.component import PackageEntryType, RemotePackageRegistry
1215
from dagster_dg.config import normalize_cli_config
1316
from dagster_dg.context import DgContext
1417
from dagster_dg.defs import (
@@ -63,42 +66,89 @@ def list_component_command(**global_options: object) -> None:
6366

6467

6568
# ########################
66-
# ##### COMPONENT TYPE
69+
# ##### PACKAGE ENTRY
6770
# ########################
6871

6972

70-
def _list_package_entries(
71-
registry: RemotePackageRegistry, output_json: bool, entry_type: Optional[str]
73+
ENTRY_TYPE_COLOR_MAP = {"component": "deep_sky_blue3", "scaffold-target": "khaki1"}
74+
75+
76+
def _pretty_entry_types(entry: PackageEntrySnap) -> Text:
77+
text = Text()
78+
for entry_type in entry.types:
79+
if len(text) > 0:
80+
text += Text(", ")
81+
text += Text(entry_type, style=ENTRY_TYPE_COLOR_MAP.get(entry_type, ""))
82+
text = Text("[") + text + Text("]")
83+
return text
84+
85+
86+
def _package_entry_table(entries: Sequence[PackageEntrySnap]) -> Table:
87+
sorted_entries = sorted(entries, key=lambda x: x.key.to_typename())
88+
table = Table(border_style="dim", show_lines=True)
89+
table.add_column("Key", style="bold cyan", no_wrap=True)
90+
table.add_column("Summary")
91+
table.add_column("Types")
92+
for entry in sorted_entries:
93+
table.add_row(entry.key.to_typename(), entry.summary, _pretty_entry_types(entry))
94+
return table
95+
96+
97+
def _all_packages_entry_table(
98+
registry: RemotePackageRegistry, name_only: bool, entry_type: Optional[PackageEntryType]
99+
) -> Table:
100+
table = Table(border_style="dim")
101+
102+
table.add_column("Package", style="bold")
103+
if not name_only:
104+
table.add_column("Entries", style="bold")
105+
106+
for package in sorted(registry.packages):
107+
if not name_only:
108+
entries = registry.get_entries(package, entry_type)
109+
inner_table = _package_entry_table(entries)
110+
table.add_row(package, inner_table)
111+
else:
112+
table.add_row(package)
113+
return table
114+
115+
116+
@list_group.command(name="packages", cls=DgClickCommand)
117+
@click.option(
118+
"--name-only",
119+
is_flag=True,
120+
default=False,
121+
help="Only display the names of the packages.",
122+
)
123+
@click.option(
124+
"--package",
125+
"-p",
126+
help="Filter by package name.",
127+
)
128+
@click.option(
129+
"--entry-type",
130+
"-t",
131+
type=click.Choice(["component", "scaffold-target"]),
132+
help="Filter by entry type.",
133+
)
134+
@dg_global_options
135+
@cli_telemetry_wrapper
136+
def list_packages_command(
137+
name_only: bool,
138+
package: Optional[str],
139+
entry_type: Optional[PackageEntryType],
140+
**global_options: object,
72141
) -> None:
73-
sorted_keys = sorted(
74-
[
75-
key
76-
for key in registry.keys()
77-
if entry_type is None or entry_type in registry.get(key).types
78-
],
79-
key=lambda k: k.to_typename(),
80-
)
81-
# JSON
82-
if output_json:
83-
output: list[dict[str, object]] = []
84-
for key in sorted_keys:
85-
obj = registry.get(key)
86-
output.append(
87-
{
88-
"key": key.to_typename(),
89-
"summary": obj.summary,
90-
}
91-
)
92-
click.echo(json.dumps(output, indent=4))
93-
# TABLE
142+
"""List registered Dagster components in the current project environment."""
143+
cli_config = normalize_cli_config(global_options, click.get_current_context())
144+
dg_context = DgContext.for_defined_registry_environment(Path.cwd(), cli_config)
145+
registry = RemotePackageRegistry.from_dg_context(dg_context)
146+
147+
if package:
148+
table = _package_entry_table(registry.get_entries(package, entry_type))
94149
else:
95-
table = Table(border_style="dim")
96-
table.add_column("Package Entry", style="bold cyan", no_wrap=True)
97-
table.add_column("Summary")
98-
for key in sorted_keys:
99-
table.add_row(key.to_typename(), registry.get(key).summary)
100-
console = Console()
101-
console.print(table)
150+
table = _all_packages_entry_table(registry, name_only, entry_type=entry_type)
151+
Console().print(table)
102152

103153

104154
@list_group.command(name="component-type", cls=DgClickCommand)
@@ -117,7 +167,15 @@ def list_component_type_command(output_json: bool, **global_options: object) ->
117167
dg_context = DgContext.for_defined_registry_environment(Path.cwd(), cli_config)
118168
registry = RemotePackageRegistry.from_dg_context(dg_context)
119169

120-
_list_package_entries(registry, output_json, "component")
170+
if output_json:
171+
output: list[dict[str, object]] = []
172+
for entry in sorted(
173+
registry.get_entries(entry_type="component"), key=lambda x: x.key.to_typename()
174+
):
175+
output.append({"key": entry.key.to_typename(), "summary": entry.summary})
176+
click.echo(json.dumps(output, indent=4))
177+
else:
178+
Console().print(_all_packages_entry_table(registry, False, "component"))
121179

122180

123181
# ########################

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

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
11
import copy
22
import json
3-
from collections.abc import Iterable, Mapping, Sequence
3+
from collections.abc import Iterable, Mapping, Sequence, Set
44
from typing import TYPE_CHECKING, Any, Optional
55

66
import dagster_shared.check as check
77
from dagster_shared.serdes import deserialize_value, serialize_value
88
from dagster_shared.serdes.objects import PackageEntryKey, PackageEntrySnap
9+
from dagster_shared.serdes.objects.package_entry import PackageEntryType
910

1011
from dagster_dg.utils import is_valid_json
1112

@@ -44,6 +45,23 @@ def __init__(self, components: dict[PackageEntryKey, PackageEntrySnap]):
4445
def empty() -> "RemotePackageRegistry":
4546
return RemotePackageRegistry({})
4647

48+
@property
49+
def packages(self) -> Set[str]:
50+
return {key.package for key in self._objects.keys()}
51+
52+
def package_entries(self, package: str) -> Set[PackageEntryKey]:
53+
return {key for key in self._objects.keys() if key.package == package}
54+
55+
def get_entries(
56+
self, package: Optional[str] = None, entry_type: Optional[PackageEntryType] = None
57+
) -> Sequence[PackageEntrySnap]:
58+
return [
59+
entry
60+
for entry in self._objects.values()
61+
if (package is None or package == entry.key.package)
62+
and (entry_type is None or entry_type in entry.types)
63+
]
64+
4765
def get(self, key: PackageEntryKey) -> PackageEntrySnap:
4866
"""Resolves a library object within the scope of a given component directory."""
4967
return self._objects[key]

‎python_modules/libraries/dagster-shared/dagster_shared/serdes/objects/package_entry.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,10 @@ class PackageEntryKey:
2222
namespace: str
2323
name: str
2424

25+
@property
26+
def package(self) -> str:
27+
return self.namespace.split(".")[0]
28+
2529
def to_typename(self) -> str:
2630
return f"{self.namespace}.{self.name}"
2731

0 commit comments

Comments
 (0)