Went down the typing Generic rabbit hole...

youtux · May 28, 2022 · c7567b6 · c7567b6
1 parent d43c7c6
commit c7567b6
Show file tree

Hide file tree

Showing 3 changed files with 118 additions and 105 deletions.
diff --git a/src/factory-stubs/base.pyi b/src/factory-stubs/base.pyi
@@ -183,26 +183,30 @@ class StubFactory(Factory[StubObject]):
     @classmethod
     def create(cls, **kwargs: Any) -> NoReturn: ...
 
-class BaseDictFactory(Generic[T], Factory[T]):
+# TODO: This should be a TMapping = TypeVar("T", bound=Mapping[KT, KV),
+#  but it doesn't seem possible with mypy
+class BaseDictFactory(Factory[T]):
     class Meta:
         abstract: bool
-    @classmethod
-    def _build(cls, model_class: Type[T], **kwargs: Any) -> T: ...  # type: ignore
-    @classmethod
-    def _create(cls, model_class: Type[T], **kwargs: Any) -> T: ...  # type: ignore
 
 class DictFactory(Generic[KT, VT], BaseDictFactory[dict[KT, VT]]):
     class Meta:
+        # This would be:
+        # model: Type[dict[KT, VT]]
+        # but mypy doesn't support it
         model: Type[dict[Any, Any]]
 
-class BaseListFactory(Generic[T], Factory[Iterable[T]]):
+# TODO: This should be a TContainer = TypeVar("T", bound=Container[KT, KV),
+#  but it doesn't seem possible with mypy
+class BaseListFactory(Generic[T], Factory[T]):
     class Meta:
         abstract: bool
-    @classmethod
-    def _build(cls, model_class: Type[Iterable[T]], **kwargs: T) -> Iterable[T]: ...  # type: ignore
 
-class ListFactory(Generic[T], BaseListFactory[T]):
+class ListFactory(Generic[T], BaseListFactory[list[T]]):
     class Meta:
+        # This would be:
+        # model: Type[list[T]]
+        # but mypy doesn't support it
         model: Type[list[Any]]
 
 TBaseFactoryType = TypeVar("TBaseFactoryType", bound=Type[BaseFactory[Any]])

diff --git a/src/factory-stubs/builder.pyi b/src/factory-stubs/builder.pyi
@@ -89,6 +89,7 @@ class StepBuilder(Generic[T]):
         self, factory_meta: base.FactoryOptions[V], extras: dict[str, Any]
     ) -> StepBuilder[V]: ...
 
+# TODO: Make this Generic[T]
 class Resolver:
     def __init__(
         self, declarations: DeclarationSet, step: BuildStep[Any], sequence: Any

diff --git a/src/factory-stubs/declarations.pyi b/src/factory-stubs/declarations.pyi
@@ -1,4 +1,5 @@
 import logging
+import typing
 from typing import (
     Any,
     Callable,
@@ -12,152 +13,158 @@ from typing import (
     TypeVar,
 )
 
-from factory.builder import BuildStep, Resolver
-
-from . import base, utils
+from . import base, utils, builder
 
 T = TypeVar("T")
+U = TypeVar("U")
 V = TypeVar("V")
 KT = TypeVar("KT")
 VT = TypeVar("VT")
 
-TFactory = TypeVar("TFactory", bound=base.Factory)
-
 logger: logging.Logger
 
-class BaseDeclaration(utils.OrderedBase):
+class BaseDeclaration(Generic[T, V], utils.OrderedBase):
     FACTORY_BUILDER_PHASE: Literal["attributes", "post_instance"]
     UNROLL_CONTEXT_BEFORE_EVALUATION: bool
     def __init__(self, **defaults: Any) -> None: ...
     def unroll_context(
-        self, instance: Resolver, step: BuildStep, context: Mapping[str, Any]
+        self,
+        instance: builder.Resolver,
+        step: builder.BuildStep[T],
+        context: Mapping[str, Any],
     ) -> dict[str, Any]: ...
     def evaluate_pre(
-        self, instance: Resolver, step: BuildStep, overrides: dict[str, Any]
-    ): ...
+        self,
+        instance: builder.Resolver,
+        step: builder.BuildStep[T],
+        overrides: dict[str, Any],
+    ) -> V: ...
     def evaluate(
-        self, instance: Resolver, step: BuildStep, extra: dict[str, Any]
-    ) -> Any: ...
+        self,
+        instance: builder.Resolver,
+        step: builder.BuildStep[T],
+        extra: dict[str, Any],
+    ) -> V: ...
 
-class OrderedDeclaration(BaseDeclaration): ...
+class OrderedDeclaration(BaseDeclaration[T, V]): ...
 
-class LazyFunction(Generic[T], BaseDeclaration):
-    function: Callable[[], T]
-    def __init__(self, function: Callable[[], T]) -> None: ...
-    def evaluate(
-        self, instance: Resolver, step: BuildStep, extra: dict[str, Any]
-    ) -> T: ...
+class LazyFunction(BaseDeclaration[T, V]):
+    function: Callable[[], V]
+    def __init__(self, function: Callable[[], V]) -> None: ...
 
-class LazyAttribute(Generic[T], BaseDeclaration):
-    function: Callable[[Resolver], T]
-    def __init__(self, function: Callable[[Resolver], T]) -> None: ...
-    def evaluate(
-        self, instance: Resolver, step: BuildStep, extra: dict[str, Any]
-    ) -> T: ...
+class LazyAttribute(BaseDeclaration[T, V]):
+    function: Callable[[builder.Resolver], V]
+    def __init__(self, function: Callable[[builder.Resolver], V]) -> None: ...
 
 class _UNSPECIFIED: ...
 
 def deepgetattr(obj: Any, name: str, default: _UNSPECIFIED | Any = ...) -> Any: ...
 
-class SelfAttribute(BaseDeclaration):
+class SelfAttribute(BaseDeclaration[T, V]):
     depth: int
     attribute_name: str
     default: _UNSPECIFIED | Any
     def __init__(
         self, attribute_name: str, default: _UNSPECIFIED | Any = ...
     ) -> None: ...
 
-class Iterator(Generic[T, V], BaseDeclaration):
-    getter: Callable[[T], V] | None
-    iterator: utils.ResetableIterator[Iterator[T]] | None
-    iterator_builder: Callable[[], utils.ResetableIterator[T]]
+class Iterator(Generic[U, T, V], BaseDeclaration[T, V]):
+    getter: Callable[[U], V] | None
+    iterator: utils.ResetableIterator[typing.Iterator[U]] | None
+    iterator_builder: Callable[[], utils.ResetableIterator[U]]
     def __init__(
         self,
-        iterator: Iterator[T],
+        iterator: typing.Iterator[U],
         cycle: bool = ...,
-        getter: Callable[[T], V] | None = ...,
+        getter: Callable[[U], V] | None = ...,
     ): ...
-    def evaluate(
-        self, instance: Resolver, step: BuildStep, extra: dict[str, Any]
-    ) -> V: ...
     def reset(self) -> None: ...
 
-class Sequence(Generic[T], BaseDeclaration):
-    function: Callable[[int], T]
-    def __init__(self, function: Callable[[int], T]) -> None: ...
-    def evaluate(
-        self, instance: Resolver, step: BuildStep, extra: dict[str, Any]
-    ) -> T: ...
+class Sequence(BaseDeclaration[T, V]):
+    # Workaround for mypy bug https://github.com/python/mypy/issues/708
+    # Otherwise it would just be this:
+    # function: Callable[[int], V]
+    @staticmethod
+    def function(sequence: int, /) -> V: ...
 
-class LazyAttributeSequence(Sequence[T]):
-    function: Callable[[Resolver, int], T]
-    def __init__(self, function: Callable[[Resolver, int], T]) -> None: ...
-    def evaluate(
-        self, instance: Resolver, step: BuildStep, extra: dict[str, Any]
-    ) -> T: ...
+    def __init__(self, function: Callable[[int], V]) -> None: ...
+
+class LazyAttributeSequence(Sequence[T, V]):
+    # Workaround for mypy bug https://github.com/python/mypy/issues/708
+    # Otherwise it would just be this:
+    # function: Callable[[builder.Resolver, int], V]
+    @staticmethod
+    def function(instance: builder.Resolver, sequence: int, /) -> V: ...  # type: ignore[override]
 
-class ContainerAttribute(Generic[T], BaseDeclaration):
-    function: Callable[[Resolver, Tuple[Resolver, ...]], T]
+    def __init__(self, function: Callable[[builder.Resolver, int], V]) -> None: ...
+
+class ContainerAttribute(BaseDeclaration[T, V]):
+    function: Callable[[builder.Resolver, Tuple[builder.Resolver, ...]], V]
     strict: bool
     def __init__(
         self,
-        function: Callable[[Resolver, Tuple[Resolver, ...]], T],
+        function: Callable[[builder.Resolver, Tuple[builder.Resolver, ...]], V],
         strict: bool = ...,
     ) -> None: ...
-    def evaluate(
-        self, instance: Resolver, step: BuildStep, extra: dict[str, Any]
-    ) -> T: ...
 
-class ParameteredAttribute(BaseDeclaration):
-    def generate(self, step: BuildStep, params: dict[str, Any]) -> Any: ...
+class ParameteredAttribute(BaseDeclaration[T, V]):
+    def generate(self, step: builder.BuildStep[T], params: dict[str, Any]) -> V: ...
 
-class _FactoryWrapper(Generic[TFactory]):
-    factory: Type[TFactory] | None
+class _FactoryWrapper(Generic[T]):
+    factory: Type[base.Factory[T]] | None
     module: str
-    def __init__(self, factory_or_path: str | Type[TFactory]) -> None: ...
-    def get(self) -> Type[TFactory]: ...
+    def __init__(self, factory_or_path: str | Type[base.Factory[T]]) -> None: ...
+    def get(self) -> Type[base.Factory[T]]: ...
 
-class SubFactory(Generic[TFactory], BaseDeclaration):
+class SubFactory(BaseDeclaration[T, V]):
     FORCE_SEQUENCE: bool
     UNROLL_CONTEXT_BEFORE_EVALUATION: bool
-    factory_wrapper: _FactoryWrapper[TFactory]
-    def __init__(self, factory: str | Type[TFactory], **kwargs: Any) -> None: ...
-    def get_factory(self) -> Type[TFactory]: ...
+    factory_wrapper: _FactoryWrapper[base.Factory[V]]
+    def __init__(self, factory: str | Type[base.Factory[V]], **kwargs: Any) -> None: ...
+    # TODO: Can evaluate(instance) be a T, or can it be only a Resolver?
+    def get_factory(self) -> Type[base.Factory[V]]: ...
 
-class Dict(Mapping[KT, VT], SubFactory[base.DictFactory[KT, VT]]):
+class Dict(Generic[T, KT, VT], SubFactory[T, base.DictFactory[KT, VT]]):
     FORCE_SEQUENCE: bool
+
+    # TODO: I'm not 100% about the type of params
     def __init__(
         self,
-        params: Mapping[KT, VT] | Iterable[Tuple[KT, VT]],
-        dict_factory: str | Type[base.Factory] = ...,
+        params: Mapping[KT, VT | SelfAttribute[Mapping[KT, VT], VT]],
+        dict_factory: str | Type[base.DictFactory[KT, VT]] = ...,
     ) -> None: ...
 
-class List(Generic[T], SubFactory[base.ListFactory[T]]):
+class List(Generic[T, V], SubFactory[T, base.ListFactory[V]]):
     FORCE_SEQUENCE: bool
+    # TODO: I'm not 100% about the type of params
     def __init__(
-        self, params: Iterable[T], list_factory: str | Type[base.Factory] = ...
+        self, params: Iterable[V | SelfAttribute[list[V], V]], list_factory: str | Type[base.ListFactory[V]] = ...
     ) -> None: ...
 
 class Skip:
     def __bool__(self) -> bool: ...
 
 SKIP: Skip
 
-class Maybe(BaseDeclaration):
-    decider: SelfAttribute
-    yes: Any
-    no: Any
+class Maybe(BaseDeclaration[T, Any]):
+    decider: SelfAttribute[T, Any]
+    # TODO: These should be Skip | U | BaseDeclaration[T, U]
+    yes: Skip | Any
+    no: Skip | Any
     def __init__(
         self,
-        decider: SelfAttribute | str,
+        decider: SelfAttribute[T, Any] | str,
         yes_declaration: Skip | Any = ...,
         no_declaration: Skip | Any = ...,
     ) -> None: ...
     def evaluate_post(
-        self, instance: Any, step: BuildStep, overrides: dict[str, Any]
+        self, instance: T, step: builder.BuildStep[T], overrides: dict[str, Any]
     ) -> Any: ...
     def evaluate_pre(
-        self, instance: Resolver, step: BuildStep, overrides: dict[str, Any]
+        self,
+        instance: builder.Resolver,
+        step: builder.BuildStep[T],
+        overrides: dict[str, Any],
     ) -> Any: ...
 
 class Parameter(utils.OrderedBase):
@@ -166,6 +173,7 @@ class Parameter(utils.OrderedBase):
     ) -> dict[str, Any]: ...
     def get_revdeps(self, parameters: dict[str, Any]) -> list[str]: ...
 
+# TODO: This requires more typevars, remove the Anys
 class SimpleParameter(Generic[T], Parameter):
     value: T
     def __init__(self, value: T) -> None: ...
@@ -175,62 +183,62 @@ class SimpleParameter(Generic[T], Parameter):
     @classmethod
     def wrap(cls, value: utils.OrderedBase | Parameter | Any) -> utils.OrderedBase: ...
 
+# TODO: this needs more typevars, remove the Anys
 class Trait(Parameter):
     overrides: dict[str, Any]
     def __init__(self, **overrides: Any) -> None: ...
     def as_declarations(
         self, field_name: str, declarations: dict[str, Any]
-    ) -> dict[str, Maybe]: ...
+    ) -> dict[str, Maybe[Any]]: ...
 
 class PostGenerationContext(NamedTuple):
     value_provided: bool
     value: Any
     extra: dict[str, Any]
 
-class PostGenerationDeclaration(BaseDeclaration):
+class PostGenerationDeclaration(BaseDeclaration[T, V]):
     def evaluate_post(
-        self, instance: Any, step: BuildStep, overrides: dict[str, Any]
-    ) -> Any: ...
-    def call(
-        self, instance: Any, step: BuildStep, context: PostGenerationContext
-    ) -> Any: ...
-
-class PostGeneration(Generic[T, V], PostGenerationDeclaration):
-    function: Callable[[T, bool, Any, ...], V]
-    def __init__(self, function: Callable[[T, bool, Any, ...], V]) -> None: ...
+        self, instance: T, step: builder.BuildStep[T], overrides: dict[str, Any]
+    ) -> V: ...
     def call(
-        self, instance: T, step: BuildStep, context: PostGenerationContext
+        self, instance: T, step: builder.BuildStep[T], context: PostGenerationContext
     ) -> V: ...
 
-class RelatedFactory(Generic[TFactory], PostGenerationDeclaration):
+
+class PostGeneration(PostGenerationDeclaration[T, V]):
+    # Workaround for mypy bug https://github.com/python/mypy/issues/708
+    # Otherwise it would just be this:
+    # function: Callable[[T, bool, Any, ...], V]
+    @staticmethod
+    def function(obj: T, create: bool, extracted: Any, /, **kwargs: Any) -> V: ...
+    def __init__(self, function: Callable[[T, bool, Any], V]) -> None: ...
+
+class RelatedFactory(PostGenerationDeclaration[T, V]):
     UNROLL_CONTEXT_BEFORE_EVALUATION: bool
     name: str
     defaults: dict[str, Any]
-    factory_wrapper: _FactoryWrapper[TFactory]
+    factory_wrapper: _FactoryWrapper[base.Factory[V]]
     def __init__(
         self,
-        factory: str | Type[TFactory],
+        factory: str | Type[base.Factory[V]],
         factory_related_name: str = ...,
         **defaults: Any,
     ) -> None: ...
-    def get_factory(self) -> Type[TFactory]: ...
+    def get_factory(self) -> Type[base.Factory[V]]: ...
 
-class RelatedFactoryList(Generic[TFactory], RelatedFactory[TFactory]):
+class RelatedFactoryList(Generic[T, V], RelatedFactory[T, list[V]]):
     size: int | Callable[[], int]
     def __init__(
         self,
-        factory: str | Type[TFactory],
+        factory: str | Type[base.Factory[V]],
         factory_related_name: str = ...,
         size: int | Callable[[], int] = ...,
         **defaults: Any,
     ) -> None: ...
-    def call(
-        self, instance: Any, step: BuildStep, context: PostGenerationContext
-    ) -> list[Any]: ...
 
 class NotProvided: ...
 
-class PostGenerationMethodCall(PostGenerationDeclaration):
+class PostGenerationMethodCall(PostGenerationDeclaration[T, V]):
     method_name: str
     method_arg: Any
     method_kwargs: dict[str, Any]