diff --git a/mypy/stubgen.py b/mypy/stubgen.py index 7ea9fa2b3b8b2..d26d731df1b28 100755 --- a/mypy/stubgen.py +++ b/mypy/stubgen.py @@ -229,6 +229,7 @@ def __init__( verbose: bool, quiet: bool, export_less: bool, + allow_class_level_aliases: bool, ) -> None: # See parse_options for descriptions of the flags. self.pyversion = pyversion @@ -247,6 +248,7 @@ def __init__( self.verbose = verbose self.quiet = quiet self.export_less = export_less + self.allow_class_level_aliases = allow_class_level_aliases class StubSource: @@ -594,6 +596,7 @@ def __init__( include_private: bool = False, analyzed: bool = False, export_less: bool = False, + allow_class_level_aliases: bool = False, ) -> None: # Best known value of __all__. self._all_ = _all_ @@ -608,6 +611,7 @@ def __init__( self._state = EMPTY self._toplevel_names: list[str] = [] self._include_private = include_private + self.allow_class_level_aliases = allow_class_level_aliases self.import_tracker = ImportTracker() # Was the tree semantically analysed before? self.analyzed = analyzed @@ -995,7 +999,8 @@ def visit_assignment_stmt(self, o: AssignmentStmt) -> None: self.process_namedtuple(lvalue, o.rvalue) continue if ( - isinstance(lvalue, NameExpr) + (self.is_top_level() or self.allow_class_level_aliases) + and isinstance(lvalue, NameExpr) and not self.is_private_name(lvalue.name) and # it is never an alias with explicit annotation @@ -1623,6 +1628,7 @@ def generate_stub_from_ast( parse_only: bool = False, include_private: bool = False, export_less: bool = False, + allow_class_level_aliases: bool = False, ) -> None: """Use analysed (or just parsed) AST to generate type stub for single file. @@ -1632,6 +1638,7 @@ def generate_stub_from_ast( gen = StubGenerator( mod.runtime_all, include_private=include_private, + allow_class_level_aliases=allow_class_level_aliases, analyzed=not parse_only, export_less=export_less, ) @@ -1695,7 +1702,12 @@ def generate_stubs(options: Options) -> None: files.append(target) with generate_guarded(mod.module, target, options.ignore_errors, options.verbose): generate_stub_from_ast( - mod, target, options.parse_only, options.include_private, options.export_less + mod, + target, + options.parse_only, + options.include_private, + options.export_less, + options.allow_class_level_aliases, ) # Separately analyse C modules using different logic. @@ -1755,6 +1767,13 @@ def parse_options(args: list[str]) -> Options: help="generate stubs for objects and members considered private " "(single leading underscore and no trailing underscores)", ) + parser.add_argument( + "--allow-class-level-aliases", + action="store_true", + help="by default variables which refer to the values of other variables " + "are only respected at the module-level. this allows aliases at the " + "class level as well.", + ) parser.add_argument( "--export-less", action="store_true", @@ -1840,6 +1859,7 @@ def parse_options(args: list[str]) -> Options: verbose=ns.verbose, quiet=ns.quiet, export_less=ns.export_less, + allow_class_level_aliases=ns.allow_class_level_aliases, ) diff --git a/test-data/unit/stubgen.test b/test-data/unit/stubgen.test index 2b1a77a1cfddb..14faa80156b6f 100644 --- a/test-data/unit/stubgen.test +++ b/test-data/unit/stubgen.test @@ -947,19 +947,18 @@ from typing import Any alias = Container[Any] -[case testAliasExceptions] -noalias1 = None -noalias2 = ... -noalias3 = True +[case testAliasOnlyToplevel] +class Foo: + alias = str [out] from _typeshed import Incomplete -noalias1: Incomplete -noalias2: Incomplete -noalias3: bool +class Foo: + alias: Incomplete [case testComplexAlias] +# flags: --allow-class-level-aliases # modules: main a from a import valid @@ -1009,6 +1008,18 @@ class A: # a.pyi valid: list[int] +[case testAliasExceptions] +noalias1 = None +noalias2 = ... +noalias3 = True + +[out] +from _typeshed import Incomplete + +noalias1: Incomplete +noalias2: Incomplete +noalias3: bool + -- More features/fixes: -- do not export deleted names