diff --git a/python_modules/dagster/dagster/_config/structured_config.py b/python_modules/dagster/dagster/_config/structured_config.py index 181575b917199..682f1e4278bb8 100644 --- a/python_modules/dagster/dagster/_config/structured_config.py +++ b/python_modules/dagster/dagster/_config/structured_config.py @@ -131,17 +131,29 @@ class _Temp(Generic[T]): @dataclass_transform() class BaseResourceMeta(pydantic.main.ModelMetaclass): + """ + Custom metaclass for Resource and PartialResource. This metaclass is responsible for + transforming the type annotations on the class to allow for partially configured resources. + """ + def __new__(self, name, bases, namespaces, **kwargs): + # Gather all type annotations from the class and its base classes annotations = namespaces.get("__annotations__", {}) for base in bases: if hasattr(base, "__annotations__"): annotations.update(base.__annotations__) for field in annotations: if not field.startswith("__"): + # Check if the annotation is a ResourceDependency if get_origin(annotations[field]) == _ResourceDep: arg = get_args(annotations[field])[0] + # If so, we treat it as a Union of a PartialResource and a Resource + # for Pydantic's sake. annotations[field] = Union[_PartialResource[arg], _Resource[arg]] elif _safe_is_subclass(annotations[field], _Resource): + # If the annotation is a Resource, we treat it as a Union of a PartialResource + # and a Resource for Pydantic's sake, so that a user can pass in a partially + # configured resource. base = annotations[field] annotations[field] = Union[_PartialResource[base], base] diff --git a/python_modules/dagster/dagster_tests/core_tests/resource_tests/test_struct_config_resources.py b/python_modules/dagster/dagster_tests/core_tests/resource_tests/test_struct_config_resources.py index f1c0f4b2308b2..f4e77c2fedf47 100644 --- a/python_modules/dagster/dagster_tests/core_tests/resource_tests/test_struct_config_resources.py +++ b/python_modules/dagster/dagster_tests/core_tests/resource_tests/test_struct_config_resources.py @@ -139,30 +139,30 @@ def another_asset(greeting: GreetingResource, hello_world_asset): def test_abc_resource(): out_txt = [] - class WriterResource(Resource, ABC): + class Writer(Resource, ABC): @abstractmethod def output(self, text: str) -> None: pass - class PrefixedWriterResource(WriterResource): + class PrefixedWriterResource(Writer): prefix: str def output(self, text: str) -> None: out_txt.append(f"{self.prefix}{text}") - class RepetitiveWriterResource(WriterResource): + class RepetitiveWriterResource(Writer): repetitions: int def output(self, text: str) -> None: out_txt.append(f"{text} " * self.repetitions) @op - def hello_world_op(writer: WriterResource): + def hello_world_op(writer: Writer): writer.output("hello, world!") # Can't instantiate abstract class with pytest.raises(TypeError): - WriterResource() # pylint: disable=abstract-class-instantiated + Writer() # pylint: disable=abstract-class-instantiated @job(resource_defs={"writer": PrefixedWriterResource(prefix="greeting: ")}) def prefixed_job(): @@ -356,18 +356,25 @@ def hello_world_asset(writer: WriterResource): def test_nested_resources(): out_txt = [] - class WriterResource(Resource): + class Writer(Resource, ABC): + @abstractmethod + def output(self, text: str) -> None: + pass + + class WriterResource(Writer): def output(self, text: str) -> None: out_txt.append(text) - class PrefixedWriterResource(WriterResource, Resource): + class PrefixedWriterResource(Writer): prefix: str def output(self, text: str) -> None: out_txt.append(f"{self.prefix}{text}") - class JsonWriterResource(WriterResource): - base_writer: WriterResource + class JsonWriterResource( + Writer, + ): + base_writer: Writer indent: int def output(self, obj: Any) -> None: