feat(mypy_plugin): item extracted results analyzed as typeddict (#87)

* feat(mypy_plugin): item extracted results analyzed as typeddict
* fix: compat issue

Author: linw1995

data_extractor/contrib/mypy/__init__.py

@@ -1,5 +1,8 @@
 # Standard Library
-from typing import Callable, Dict, Optional, Type
+import logging
+
+from functools import partial
+from typing import Callable, Dict, List, Optional, Type
 
 # Third Party Library
 from mypy.checker import TypeChecker
@@ -15,11 +18,21 @@
     TypeInfo,
 )
 from mypy.nodes import Var as VarExpr
-from mypy.plugin import FunctionContext, MethodSigContext, Plugin
+from mypy.options import Options
+from mypy.plugin import (
+    DynamicClassDefContext,
+    FunctionContext,
+    MethodSigContext,
+    Plugin,
+)
+from mypy.semanal import SemanticAnalyzerInterface
+from mypy.semanal_typeddict import TypedDictAnalyzer
 from mypy.traverser import TraverserVisitor
 from mypy.types import AnyType, CallableType, Instance
 from mypy.types import Type as MypyType
-from mypy.types import TypeOfAny, UninhabitedType, UnionType
+from mypy.types import TypedDictType, TypeOfAny, UninhabitedType, UnionType
+
+logger = logging.getLogger(__name__)
 
 
 class RelationshipVisitor(TraverserVisitor):
@@ -81,7 +94,13 @@ def visit_assignment_stmt(self, o):
 
 
 class DataExtractorPlugin(Plugin):
-    cache: Dict[str, Dict[str, str]] = {}
+    cache: Dict[str, Dict[str, str]]
+    item_typeddict_mapping: Dict[str, TypedDictType]
+
+    def __init__(self, options: Options) -> None:
+        super().__init__(options)
+        self.cache = {}
+        self.item_typeddict_mapping = {}
 
     def get_current_code(self, ctx: FunctionContext) -> MypyFile:
         api = ctx.api
@@ -112,12 +131,6 @@ def check_field_generic_type(self, ctx: FunctionContext) -> MypyType:
         if rv_type.args and not isinstance(rv_type.args[0], UninhabitedType):
             return rv_type
 
-        # # check parameter "type"
-        # idx = ctx.callee_arg_names.index("type")
-        # arg = ctx.args[type_idx]
-        # if arg:
-        #    args = [arg[0].node]
-        # else:
         if not self.options.disallow_any_generics:
             return self.apply_any_generic(type=rv_type)
         else:
@@ -139,11 +152,12 @@ def check_is_many(self, ctx: FunctionContext) -> bool:
 
         return False
 
-    def prepare_type_annotations(self, ctx: FunctionContext) -> MypyType:
+    def prepare_type_annotations(self, ctx: FunctionContext, fullname: str) -> MypyType:
+        logger.debug("prepare_type_annotations %r", fullname)
+
         # check parameter "is_many"
         expr = ctx.context
         assert isinstance(expr, CallExpr)
-
         relationship = self.anal_code(self.get_current_code(ctx))
         lvalue_key = str((expr.line, expr.column))
         if lvalue_key in relationship:
@@ -171,24 +185,30 @@ def prepare_type_annotations(self, ctx: FunctionContext) -> MypyType:
         rv_type = self.check_field_generic_type(ctx)
         return rv_type
 
-    def is_extractor_cls(self, fullname: str) -> bool:
+    def is_extractor_cls(self, fullname: str, is_item_subcls=False) -> bool:
         node = self.lookup_fully_qualified(fullname)
         if node is not None:
             typenode = node.node
             if isinstance(typenode, TypeInfo):
-                return typenode.has_base("data_extractor.item.Field")
+                if is_item_subcls:
+                    return typenode.has_base("data_extractor.item.Item")
+                else:
+                    return typenode.has_base("data_extractor.item.Field")
 
         return False
 
     def get_function_hook(
         self, fullname: str
     ) -> Optional[Callable[[FunctionContext], MypyType]]:
+        logger.debug("get_function_hook %r", fullname)
         if self.is_extractor_cls(fullname):
-            return self.prepare_type_annotations
+            return partial(self.prepare_type_annotations, fullname=fullname)
 
         return super().get_function_hook(fullname)
 
-    def apply_is_many_on_extract_method(self, ctx: MethodSigContext) -> CallableType:
+    def apply_is_many_on_extract_method(
+        self, ctx: MethodSigContext, fullname: str
+    ) -> CallableType:
         origin: CallableType = ctx.default_signature
         origin_ret_type = origin.ret_type
         assert isinstance(origin_ret_type, UnionType)
@@ -211,6 +231,7 @@ def apply_is_many_on_extract_method(self, ctx: MethodSigContext) -> CallableType
             key = str((obj.line, obj.column))
             # metadata = obj.type.type.metadata
 
+        logger.debug("apply_is_many %r %r %r", fullname, key, metadata)
         if key in metadata:
             is_many = metadata[key]["is_many"]
             ret_type = origin_ret_type.items[int(is_many)]
@@ -227,14 +248,108 @@ def is_extract_method(self, fullname: str) -> bool:
             return self.is_extractor_cls(fullname[: -len(suffix)])
         return False
 
+    def apply_extract_method(
+        self, ctx: MethodSigContext, fullname: str
+    ) -> CallableType:
+        rv = self.apply_is_many_on_extract_method(ctx, fullname)
+
+        # apply item typeddict
+        item_classname = fullname[: -len(".extract")]
+        if item_classname in self.item_typeddict_mapping:
+            logger.debug("apply_extract_method %r %r", fullname, rv.ret_type)
+            original = rv.ret_type
+            typeddict = self.item_typeddict_mapping[item_classname]
+            ret_type: Optional[MypyType]
+            if isinstance(original, AnyType):
+                ret_type = typeddict
+            else:
+                assert isinstance(original, Instance)
+                if original.type.name == "list":
+                    ret_type = original
+                    ret_type.args = (typeddict,)
+                else:
+                    api = ctx.api
+                    assert isinstance(api, TypeChecker)
+                    api.fail(
+                        "Cant determine extract method return type", context=ctx.context
+                    )
+                    ret_type = None
+
+            if ret_type is not None:
+                rv = rv.copy_modified(ret_type=ret_type)
+
+        logger.debug(
+            "apply_extract_method %r %r %r", fullname, rv, self.item_typeddict_mapping
+        )
+        return rv
+
     def get_method_signature_hook(
         self, fullname: str
     ) -> Optional[Callable[[MethodSigContext], CallableType]]:
         if self.is_extract_method(fullname):
-            return self.apply_is_many_on_extract_method
-
+            return partial(self.apply_extract_method, fullname=fullname)
         return super().get_method_signature_hook(fullname)
 
+    def prepare_typeddict(self, ctx: DynamicClassDefContext, fullname: str) -> None:
+        logger.debug("prepare_typeddict %r", fullname)
+        if fullname in self.item_typeddict_mapping:
+            return
+
+        api = ctx.api
+        assert isinstance(api, SemanticAnalyzerInterface)
+        analyzer = TypedDictAnalyzer(api.options, api, api.msg)  # type: ignore
+
+        items: List[str] = []
+        types: List[MypyType] = []
+        callee = ctx.call.callee
+        assert isinstance(callee, NameExpr)
+        node = callee.node
+        assert isinstance(node, TypeInfo)
+        for block in node.defn.defs.body:
+            if not isinstance(block, AssignmentStmt):
+                continue
+
+            rvalue = block.rvalue
+            if not isinstance(rvalue, CallExpr):
+                continue
+
+            rvalue_type: MypyType
+            callee = rvalue.callee
+            if isinstance(callee, IndexExpr):
+                index = callee.index
+                assert isinstance(index, NameExpr)
+                name = index.fullname
+                assert name is not None
+                named_type = api.named_type_or_none(name, [])
+                assert named_type is not None
+                rvalue_type = named_type
+            else:
+                rvalue_type = AnyType(TypeOfAny.special_form)
+
+            for lvalue in block.lvalues:
+                assert isinstance(lvalue, NameExpr)
+                items.append(lvalue.name)
+                types.append(rvalue_type)
+
+        callee = ctx.call.callee
+        assert isinstance(callee, NameExpr)
+        typeinfo = analyzer.build_typeddict_typeinfo(
+            callee.name, items, types, set(items), -1
+        )
+        assert typeinfo.typeddict_type is not None
+        self.item_typeddict_mapping[fullname] = typeinfo.typeddict_type
+        logger.debug("prepare_typeddict %r %r", fullname, self.item_typeddict_mapping)
+
+    def get_dynamic_class_hook(
+        self, fullname: str
+    ) -> Optional[Callable[[DynamicClassDefContext], None]]:
+        logger.debug("dynamic_class_hook %r", fullname)
+        if self.options.python_version >= (3, 8):
+            if self.is_extractor_cls(fullname, is_item_subcls=True):
+                return partial(self.prepare_typeddict, fullname=fullname)
+
+        return super().get_dynamic_class_hook(fullname)
+
 
 def plugin(version: str) -> Type[Plugin]:
     return DataExtractorPlugin

tests/typesafety/test_extracted_typed_dict.yml

@@ -0,0 +1,30 @@
+- case: item_extracted_result_is_typeddict
+  skip: sys.version_info.minor < 8
+  main: |
+    from tests.utils import D
+    from data_extractor.item import Item, Field
+
+    class Point2D(Item):
+        x = Field[int](D())
+        y = Field[int](D())
+
+    p = Point2D(D())
+    rv = p.extract({"x": 1, "y": 3})
+    reveal_type(rv)
+  out: |
+    main:10: note: Revealed type is "TypedDict({'x': builtins.int, 'y': builtins.int})"
+- case: item_extracted_many_results_are_typeddict
+  skip: sys.version_info.minor < 8
+  main: |
+    from tests.utils import D
+    from data_extractor.item import Item, Field
+
+    class Point2D(Item):
+        x = Field[int](D())
+        y = Field[int](D())
+
+    p = Point2D(D(), is_many=True)
+    rvs = p.extract([{"x": 1, "y": 3}])
+    reveal_type(rvs)
+  out: |
+    main:10: note: Revealed type is "builtins.list[TypedDict({'x': builtins.int, 'y': builtins.int})]"

tests/typesafety/test_generic.yml

@@ -133,6 +133,25 @@
     main:9: note: Revealed type is "builtins.str*"
     main:12: note: Revealed type is "builtins.list[builtins.str*]"
 - case: item_extract_with_flag_is_many
+  skip: sys.version_info.minor < 8
+  main: |
+    from tests.utils import D
+    from data_extractor import RV, Item
+
+    class C(Item[RV]):
+        pass
+
+    f1 = C(D())
+    rv = f1.extract([1])
+    reveal_type(rv)
+    f2 = C(D(), is_many=True)
+    rvs = f2.extract([1])
+    reveal_type(rvs)
+  out: |
+    main:9: note: Revealed type is "TypedDict({})"
+    main:12: note: Revealed type is "builtins.list[TypedDict({})]"
+- case: item_extract_with_flag_is_many/compat
+  skip: sys.version_info.minor >= 8
   main: |
     from tests.utils import D
     from data_extractor import RV, Item