Skip to content

Commit

Permalink
Support TypedDict functional syntax as class base type (#16703)
Browse files Browse the repository at this point in the history
Fixes #16701

This PR allows `TypedDict(...)` calls to be used as a base class. This
fixes the error emitted by mypy described in
#16701 .
  • Loading branch information
anniel-stripe committed Jan 14, 2024
1 parent 186ace3 commit b1fe23f
Show file tree
Hide file tree
Showing 3 changed files with 33 additions and 4 deletions.
8 changes: 8 additions & 0 deletions mypy/semanal.py
Original file line number Diff line number Diff line change
Expand Up @@ -2169,8 +2169,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:
Expand Down
19 changes: 15 additions & 4 deletions mypy/semanal_typeddict.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -117,7 +119,13 @@ 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, 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:
if "TypedDict" not in typeddict_bases_set:
typeddict_bases_set.add("TypedDict")
else:
Expand Down Expand Up @@ -176,19 +184,22 @@ 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] = []
else:
assert isinstance(base, IndexExpr)
elif isinstance(base, IndexExpr):
assert isinstance(base.base, RefExpr)
assert isinstance(base.base.node, TypeInfo)
info = base.base.node
args = self.analyze_base_args(base, ctx)
if args is None:
return
base_args = args
else:
assert isinstance(base, CallExpr)
assert isinstance(base.analyzed, TypedDictExpr)
info = base.analyzed.info

assert info.typeddict_type is not None
base_typed_dict = info.typeddict_type
Expand Down
10 changes: 10 additions & 0 deletions test-data/unit/check-typeddict.test
Original file line number Diff line number Diff line change
Expand Up @@ -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]

0 comments on commit b1fe23f

Please sign in to comment.