From cd83a1a67b4e193f6748c42191ffe8fa1af637e0 Mon Sep 17 00:00:00 2001 From: Annie Li Date: Fri, 22 Dec 2023 11:46:58 -0800 Subject: [PATCH 1/2] Support TypedDict functional syntax in class base type --- mypy/semanal.py | 8 ++++++++ mypy/semanal_typeddict.py | 17 ++++++++++++++--- test-data/unit/check-typeddict.test | 10 ++++++++++ 3 files changed, 32 insertions(+), 3 deletions(-) diff --git a/mypy/semanal.py b/mypy/semanal.py index 4128369ace5d..eb4039cb151a 100644 --- a/mypy/semanal.py +++ b/mypy/semanal.py @@ -2164,8 +2164,16 @@ def analyze_base_classes( if ( isinstance(base_expr, RefExpr) and base_expr.fullname in TYPED_NAMEDTUPLE_NAMES + TPDICT_NAMES + ) or ( + isinstance(base_expr, CallExpr) + and isinstance(base_expr.callee, RefExpr) + and base_expr.callee.fullname in TPDICT_NAMES ): # Ignore magic bases for now. + # For example: + # class Foo(TypedDict): ... # RefExpr + # class Foo(NamedTuple): ... # RefExpr + # class Foo(TypedDict("Foo", {"a": int})): ... # CallExpr continue try: diff --git a/mypy/semanal_typeddict.py b/mypy/semanal_typeddict.py index 13aab4de65e4..4c30d4da15fd 100644 --- a/mypy/semanal_typeddict.py +++ b/mypy/semanal_typeddict.py @@ -79,6 +79,8 @@ def analyze_typeddict_classdef(self, defn: ClassDef) -> tuple[bool, TypeInfo | N """ possible = False for base_expr in defn.base_type_exprs: + if isinstance(base_expr, CallExpr): + base_expr = base_expr.callee if isinstance(base_expr, IndexExpr): base_expr = base_expr.base if isinstance(base_expr, RefExpr): @@ -117,7 +119,12 @@ def analyze_typeddict_classdef(self, defn: ClassDef) -> tuple[bool, TypeInfo | N typeddict_bases: list[Expression] = [] typeddict_bases_set = set() for expr in defn.base_type_exprs: - if isinstance(expr, RefExpr) and expr.fullname in TPDICT_NAMES: + ok, info, _ = self.check_typeddict(expr, None, False) + if ok and info is not None: + # expr is a CallExpr + typeddict_bases_set.add(info.fullname) + typeddict_bases.append(expr) + elif isinstance(expr, RefExpr) and expr.fullname in TPDICT_NAMES: if "TypedDict" not in typeddict_bases_set: typeddict_bases_set.add("TypedDict") else: @@ -180,8 +187,7 @@ def add_keys_and_types_from_base( assert isinstance(base.node, TypeInfo) info = base.node base_args: list[Type] = [] - else: - assert isinstance(base, IndexExpr) + elif isinstance(base, IndexExpr): assert isinstance(base.base, RefExpr) assert isinstance(base.base.node, TypeInfo) info = base.base.node @@ -189,6 +195,11 @@ def add_keys_and_types_from_base( if args is None: return base_args = args + else: + assert isinstance(base, CallExpr) + assert isinstance(base.analyzed, TypedDictExpr) + info = base.analyzed.info + base_args: list[Type] = [] assert info.typeddict_type is not None base_typed_dict = info.typeddict_type diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test index d8022f85574c..625b82936e8c 100644 --- a/test-data/unit/check-typeddict.test +++ b/test-data/unit/check-typeddict.test @@ -3438,3 +3438,13 @@ class TotalInTheMiddle(TypedDict, a=1, total=True, b=2, c=3): # E: Unexpected k ... [builtins fixtures/dict.pyi] [typing fixtures/typing-typeddict.pyi] + +[case testCanCreateClassWithFunctionBasedTypedDictBase] +from mypy_extensions import TypedDict + +class Params(TypedDict("Params", {'x': int})): + pass + +p: Params = {'x': 2} +reveal_type(p) # N: Revealed type is "TypedDict('__main__.Params', {'x': builtins.int})" +[builtins fixtures/dict.pyi] From 45541625bcd2494f98978aa5f694b60b82af5854 Mon Sep 17 00:00:00 2001 From: Annie Li Date: Fri, 22 Dec 2023 12:04:32 -0800 Subject: [PATCH 2/2] Fix type errors --- mypy/semanal_typeddict.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mypy/semanal_typeddict.py b/mypy/semanal_typeddict.py index 4c30d4da15fd..8ca067b6e6ea 100644 --- a/mypy/semanal_typeddict.py +++ b/mypy/semanal_typeddict.py @@ -119,9 +119,10 @@ def analyze_typeddict_classdef(self, defn: ClassDef) -> tuple[bool, TypeInfo | N typeddict_bases: list[Expression] = [] typeddict_bases_set = set() for expr in defn.base_type_exprs: - ok, info, _ = self.check_typeddict(expr, None, False) - if ok and info is not None: + ok, maybe_type_info, _ = self.check_typeddict(expr, None, False) + if ok and maybe_type_info is not None: # expr is a CallExpr + info = maybe_type_info typeddict_bases_set.add(info.fullname) typeddict_bases.append(expr) elif isinstance(expr, RefExpr) and expr.fullname in TPDICT_NAMES: @@ -183,10 +184,10 @@ def add_keys_and_types_from_base( required_keys: set[str], ctx: Context, ) -> None: + base_args: list[Type] = [] if isinstance(base, RefExpr): assert isinstance(base.node, TypeInfo) info = base.node - base_args: list[Type] = [] elif isinstance(base, IndexExpr): assert isinstance(base.base, RefExpr) assert isinstance(base.base.node, TypeInfo) @@ -199,7 +200,6 @@ def add_keys_and_types_from_base( assert isinstance(base, CallExpr) assert isinstance(base.analyzed, TypedDictExpr) info = base.analyzed.info - base_args: list[Type] = [] assert info.typeddict_type is not None base_typed_dict = info.typeddict_type