From d6824ee104bb0bade2dd62e0f701b2e7b0e41f19 Mon Sep 17 00:00:00 2001 From: Jordandev678 <20153053+Jordandev678@users.noreply.github.com> Date: Wed, 18 Sep 2024 18:34:52 +0000 Subject: [PATCH 1/8] Add special handling for typing.get_args --- mypy/checkexpr.py | 33 +++++++++++ test-data/unit/check-get-args.test | 75 +++++++++++++++++++++++++ test-data/unit/fixtures/typing-full.pyi | 2 + 3 files changed, 110 insertions(+) create mode 100644 test-data/unit/check-get-args.test diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 22595c85e702..5150d0f4a735 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -544,6 +544,39 @@ def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) -> self.msg.cannot_use_function_with_type(e.callee.name, "TypedDict", e) elif typ.node.is_newtype: self.msg.cannot_use_function_with_type(e.callee.name, "NewType", e) + if ( + isinstance(e.callee, NameExpr) + and e.callee.fullname == "typing.get_args" + and len(e.args) == 1 + ): + #Special hanlding for get_args(), returns a typed tuple + #with the type set by the input + typ = None + if isinstance(e.args[0], IndexExpr): + self.accept(e.args[0].index) + typ = self.chk.lookup_type(e.args[0].index) + else: + try: + node = self.chk.lookup_qualified(e.args[0].name) + except KeyError: + # Undefined names should already be reported in semantic analysis. + pass + if node: + if isinstance(node.node, TypeAlias): + #Resolve type + typ = get_proper_type(node.node.target) + else: + typ = node.node.type + if ( typ is not None + and isinstance(typ, UnionType) + and all([isinstance(t, LiteralType) for t in typ.items]) + ): + # Returning strings is defined but order isn't so + # we need to return type * len of the union + return TupleType([typ] * len(typ.items), fallback=self.named_type("builtins.tuple")) + else: + # Fall back to what we did anyway (Tuple[Any]) + return TupleType([AnyType(TypeOfAny.special_form)], fallback=self.named_type("builtins.tuple")) self.try_infer_partial_type(e) type_context = None if isinstance(e.callee, LambdaExpr): diff --git a/test-data/unit/check-get-args.test b/test-data/unit/check-get-args.test new file mode 100644 index 000000000000..0d40c0a4e83b --- /dev/null +++ b/test-data/unit/check-get-args.test @@ -0,0 +1,75 @@ +[case getArgsReturnTypes] +from typing import Literal, Union, get_args +literals: Literal["a", "bc"] = "a" +unionliterals: Union[Literal["a"], Literal["bc"]] = "a" +intandliterals: Union[Literal["a", "bc"], int] = "a" +intandunionliterals: Union[Literal["a"], Literal["bc"], int] = "a" +inttype: int = 1 +reveal_type(get_args(literals)) # N: Revealed type is "Tuple[Union[Literal['a'], Literal['bc']], Union[Literal['a'], Literal['bc']]]" +reveal_type(get_args(unionliterals)) # N: Revealed type is "Tuple[Union[Literal['a'], Literal['bc']], Union[Literal['a'], Literal['bc']]]" +reveal_type(get_args(intandliterals)) # N: Revealed type is "Tuple[Any]" +reveal_type(get_args(intandunionliterals)) # N: Revealed type is "Tuple[Any]" +reveal_type(get_args(inttype)) # N: Revealed type is "Tuple[Any]" +[builtins fixtures/primitives.pyi] +[typing fixtures/typing-full.pyi] + +[case getArgsVarTypesWithNarrowing] +from typing import Literal, get_args +from typing_extensions import TypeAlias +normalImplicit = Literal["a", "bc"] +normalExplicit: TypeAlias = Literal["a", "bc"] +reveal_type(get_args(normalImplicit)) # N: Revealed type is "Tuple[Union[Literal['a'], Literal['bc']], Union[Literal['a'], Literal['bc']]]" +reveal_type(get_args(normalExplicit)) # N: Revealed type is "Tuple[Union[Literal['a'], Literal['bc']], Union[Literal['a'], Literal['bc']]]" +#reveal_type(get_args(Literal["a", "bc"])) +[builtins fixtures/primitives.pyi] +[typing fixtures/typing-full.pyi] + +[case testNarrowingInType] +from typing import Literal, get_args +type_alpha = Literal["a", "b", "c"] +strIn: str = "c" +strOut: str = "d" +if strIn in get_args(type_alpha): + reveal_type(strIn) # N: Revealed type is "Union[Literal['a'], Literal['b'], Literal['c']]" +else: + reveal_type(strIn) # N: Revealed type is "builtins.str" +if strOut in get_args(type_alpha): + reveal_type(strOut) # N: Revealed type is "Union[Literal['a'], Literal['b'], Literal['c']]" +else: + reveal_type(strOut) # N: Revealed type is "builtins.str" +[builtins fixtures/primitives.pyi] +[typing fixtures/typing-full.pyi] + +[case testNarrowingNotInType] +from typing import Literal, get_args +type_alpha = Literal["a", "b", "c"] +strIn: str = "c" +strOut: str = "d" +if strIn not in get_args(type_alpha): + reveal_type(strIn) # N: Revealed type is "builtins.str" +else: + reveal_type(strIn) # N: Revealed type is "Union[Literal['a'], Literal['b'], Literal['c']]" +if strOut not in get_args(type_alpha): + reveal_type(strOut) # N: Revealed type is "builtins.str" +else: + reveal_type(strOut) # N: Revealed type is "Union[Literal['a'], Literal['b'], Literal['c']]" +[builtins fixtures/primitives.pyi] +[typing fixtures/typing-full.pyi] + +[case i15106] +from typing import Literal, get_args, Optional +from typing_extensions import TypeAlias +ExpectedUserInput: TypeAlias = Literal[ + "these", "strings", "are", "expected", "user", "input" +] +def external_function(input: str) -> Optional[str]: + if input not in get_args(ExpectedUserInput): + reveal_type(input) # N: Revealed type is "builtins.str" + return None + reveal_type(input) # N: Revealed type is "Union[Literal['these'], Literal['strings'], Literal['are'], Literal['expected'], Literal['user'], Literal['input']]" + return _internal_function(input) + +def _internal_function(input: ExpectedUserInput) -> str: + return "User input: {input}" +[builtins fixtures/primitives.pyi] +[typing fixtures/typing-full.pyi] \ No newline at end of file diff --git a/test-data/unit/fixtures/typing-full.pyi b/test-data/unit/fixtures/typing-full.pyi index 8e0116aab1c2..e732d532a537 100644 --- a/test-data/unit/fixtures/typing-full.pyi +++ b/test-data/unit/fixtures/typing-full.pyi @@ -204,6 +204,8 @@ def dataclass_transform( ) -> Callable[[T], T]: ... def override(__arg: T) -> T: ... +def get_args(tp: T) -> T: ... + # Was added in 3.11 def reveal_type(__obj: T) -> T: ... From e4e6dc556bc25dfd82acd34909f65174fed8ac5a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 18 Sep 2024 18:51:24 +0000 Subject: [PATCH 2/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mypy/checkexpr.py | 21 +++++++++++++-------- test-data/unit/check-get-args.test | 2 +- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 5150d0f4a735..583e05c0de9e 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -549,8 +549,8 @@ def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) -> and e.callee.fullname == "typing.get_args" and len(e.args) == 1 ): - #Special hanlding for get_args(), returns a typed tuple - #with the type set by the input + # Special hanlding for get_args(), returns a typed tuple + # with the type set by the input typ = None if isinstance(e.args[0], IndexExpr): self.accept(e.args[0].index) @@ -559,24 +559,29 @@ def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) -> try: node = self.chk.lookup_qualified(e.args[0].name) except KeyError: - # Undefined names should already be reported in semantic analysis. - pass + # Undefined names should already be reported in semantic analysis. + pass if node: if isinstance(node.node, TypeAlias): - #Resolve type + # Resolve type typ = get_proper_type(node.node.target) else: typ = node.node.type - if ( typ is not None + if ( + typ is not None and isinstance(typ, UnionType) and all([isinstance(t, LiteralType) for t in typ.items]) ): # Returning strings is defined but order isn't so # we need to return type * len of the union - return TupleType([typ] * len(typ.items), fallback=self.named_type("builtins.tuple")) + return TupleType( + [typ] * len(typ.items), fallback=self.named_type("builtins.tuple") + ) else: # Fall back to what we did anyway (Tuple[Any]) - return TupleType([AnyType(TypeOfAny.special_form)], fallback=self.named_type("builtins.tuple")) + return TupleType( + [AnyType(TypeOfAny.special_form)], fallback=self.named_type("builtins.tuple") + ) self.try_infer_partial_type(e) type_context = None if isinstance(e.callee, LambdaExpr): diff --git a/test-data/unit/check-get-args.test b/test-data/unit/check-get-args.test index 0d40c0a4e83b..9dab01f284e4 100644 --- a/test-data/unit/check-get-args.test +++ b/test-data/unit/check-get-args.test @@ -72,4 +72,4 @@ def external_function(input: str) -> Optional[str]: def _internal_function(input: ExpectedUserInput) -> str: return "User input: {input}" [builtins fixtures/primitives.pyi] -[typing fixtures/typing-full.pyi] \ No newline at end of file +[typing fixtures/typing-full.pyi] From d72f74bccc2cfdc37317414d7364023f45dd9f3a Mon Sep 17 00:00:00 2001 From: Jordandev678 <20153053+Jordandev678@users.noreply.github.com> Date: Wed, 18 Sep 2024 19:12:10 +0000 Subject: [PATCH 3/8] Fix unnecessary list comprehension --- mypy/checkexpr.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 583e05c0de9e..f24b5b6eefd8 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -570,7 +570,7 @@ def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) -> if ( typ is not None and isinstance(typ, UnionType) - and all([isinstance(t, LiteralType) for t in typ.items]) + and all(isinstance(t, LiteralType) for t in typ.items) ): # Returning strings is defined but order isn't so # we need to return type * len of the union From c8ae0e9dcba50f21dfd75db71880e72b86cda709 Mon Sep 17 00:00:00 2001 From: Jordandev678 <20153053+Jordandev678@users.noreply.github.com> Date: Wed, 18 Sep 2024 19:41:13 +0000 Subject: [PATCH 4/8] Move get_proper_type section into try block with lookup_qualified --- mypy/checkexpr.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index f24b5b6eefd8..7d5af0fba6cc 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -558,15 +558,15 @@ def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) -> else: try: node = self.chk.lookup_qualified(e.args[0].name) + if node: + if isinstance(node.node, TypeAlias): + # Resolve type + typ = get_proper_type(node.node.target) + else: + typ = node.node.type except KeyError: # Undefined names should already be reported in semantic analysis. pass - if node: - if isinstance(node.node, TypeAlias): - # Resolve type - typ = get_proper_type(node.node.target) - else: - typ = node.node.type if ( typ is not None and isinstance(typ, UnionType) From d314b67a3d1c07a871cf510de86a357c8965661c Mon Sep 17 00:00:00 2001 From: Jordandev678 <20153053+Jordandev678@users.noreply.github.com> Date: Wed, 18 Sep 2024 20:49:59 +0000 Subject: [PATCH 5/8] Update node checks to be more type specific and rename typ to argtyp --- mypy/checkexpr.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 7d5af0fba6cc..21ed7e420ff4 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -551,31 +551,33 @@ def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) -> ): # Special hanlding for get_args(), returns a typed tuple # with the type set by the input - typ = None + argtyp = None if isinstance(e.args[0], IndexExpr): self.accept(e.args[0].index) - typ = self.chk.lookup_type(e.args[0].index) - else: + argtyp = self.chk.lookup_type(e.args[0].index) + elif isinstance(e.args[0], NameExpr): try: node = self.chk.lookup_qualified(e.args[0].name) if node: if isinstance(node.node, TypeAlias): # Resolve type - typ = get_proper_type(node.node.target) - else: - typ = node.node.type + argtyp = node.node.target + elif isinstance(node.node, Var): + argtyp = node.node.type except KeyError: # Undefined names should already be reported in semantic analysis. pass + if argtyp is not None: + argtyp = get_proper_type(argtyp) if ( - typ is not None - and isinstance(typ, UnionType) - and all(isinstance(t, LiteralType) for t in typ.items) + argtyp is not None + and isinstance(argtyp, UnionType) + and all(isinstance(get_proper_type(t), LiteralType) for t in argtyp.items) ): # Returning strings is defined but order isn't so # we need to return type * len of the union return TupleType( - [typ] * len(typ.items), fallback=self.named_type("builtins.tuple") + [argtyp] * len(argtyp.items), fallback=self.named_type("builtins.tuple") ) else: # Fall back to what we did anyway (Tuple[Any]) From a01d55d44de0c9806e5814b294935dbb67e51515 Mon Sep 17 00:00:00 2001 From: Jordandev678 <20153053+Jordandev678@users.noreply.github.com> Date: Thu, 19 Sep 2024 15:36:51 +0000 Subject: [PATCH 6/8] Fall though instead of returning Tuple[Any] --- mypy/checkexpr.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 21ed7e420ff4..5482ab562c1e 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -579,11 +579,6 @@ def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) -> return TupleType( [argtyp] * len(argtyp.items), fallback=self.named_type("builtins.tuple") ) - else: - # Fall back to what we did anyway (Tuple[Any]) - return TupleType( - [AnyType(TypeOfAny.special_form)], fallback=self.named_type("builtins.tuple") - ) self.try_infer_partial_type(e) type_context = None if isinstance(e.callee, LambdaExpr): From be9c2785630cb0dd132658b5b6278fbf9b34e489 Mon Sep 17 00:00:00 2001 From: Jordandev678 <20153053+Jordandev678@users.noreply.github.com> Date: Thu, 19 Sep 2024 16:16:35 +0000 Subject: [PATCH 7/8] Update get-args test case --- test-data/unit/check-get-args.test | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test-data/unit/check-get-args.test b/test-data/unit/check-get-args.test index 9dab01f284e4..a244ee68e121 100644 --- a/test-data/unit/check-get-args.test +++ b/test-data/unit/check-get-args.test @@ -7,9 +7,9 @@ intandunionliterals: Union[Literal["a"], Literal["bc"], int] = "a" inttype: int = 1 reveal_type(get_args(literals)) # N: Revealed type is "Tuple[Union[Literal['a'], Literal['bc']], Union[Literal['a'], Literal['bc']]]" reveal_type(get_args(unionliterals)) # N: Revealed type is "Tuple[Union[Literal['a'], Literal['bc']], Union[Literal['a'], Literal['bc']]]" -reveal_type(get_args(intandliterals)) # N: Revealed type is "Tuple[Any]" -reveal_type(get_args(intandunionliterals)) # N: Revealed type is "Tuple[Any]" -reveal_type(get_args(inttype)) # N: Revealed type is "Tuple[Any]" +reveal_type(get_args(intandliterals)) # N: Revealed type is "Union[Literal['a'], Literal['bc'], builtins.int]" +reveal_type(get_args(intandunionliterals)) # N: Revealed type is "Union[Literal['a'], Literal['bc'], builtins.int]" +reveal_type(get_args(inttype)) # N: Revealed type is "builtins.int" [builtins fixtures/primitives.pyi] [typing fixtures/typing-full.pyi] From 2ed83ed7dcda2c8a01cda5049ba7dc0791955c42 Mon Sep 17 00:00:00 2001 From: Jordandev678 <20153053+Jordandev678@users.noreply.github.com> Date: Thu, 19 Sep 2024 16:32:42 +0000 Subject: [PATCH 8/8] Update case for types directly in get_args call --- test-data/unit/check-get-args.test | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test-data/unit/check-get-args.test b/test-data/unit/check-get-args.test index a244ee68e121..d6252d625f67 100644 --- a/test-data/unit/check-get-args.test +++ b/test-data/unit/check-get-args.test @@ -20,7 +20,7 @@ normalImplicit = Literal["a", "bc"] normalExplicit: TypeAlias = Literal["a", "bc"] reveal_type(get_args(normalImplicit)) # N: Revealed type is "Tuple[Union[Literal['a'], Literal['bc']], Union[Literal['a'], Literal['bc']]]" reveal_type(get_args(normalExplicit)) # N: Revealed type is "Tuple[Union[Literal['a'], Literal['bc']], Union[Literal['a'], Literal['bc']]]" -#reveal_type(get_args(Literal["a", "bc"])) +reveal_type(get_args(Literal["a", "bc"])) # N: Revealed type is "typing._SpecialForm" [builtins fixtures/primitives.pyi] [typing fixtures/typing-full.pyi]