Skip to content

Commit f18850e

Browse files
committed
stubgen: fix FunctionContext.fullname for nested classes
1 parent 7237d55 commit f18850e

File tree

4 files changed

+43
-10
lines changed

4 files changed

+43
-10
lines changed

mypy/stubgen.py

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -468,14 +468,18 @@ def __init__(
468468
self._vars: list[list[str]] = [[]]
469469
# What was generated previously in the stub file.
470470
self._state = EMPTY
471-
self._current_class: ClassDef | None = None
471+
self._class_stack: list[ClassDef] = []
472472
# Was the tree semantically analysed before?
473473
self.analyzed = analyzed
474474
# Short names of methods defined in the body of the current class
475475
self.method_names: set[str] = set()
476476
self.processing_enum = False
477477
self.processing_dataclass = False
478478

479+
@property
480+
def _current_class(self) -> ClassDef | None:
481+
return self._class_stack[-1] if self._class_stack else None
482+
479483
def visit_mypy_file(self, o: MypyFile) -> None:
480484
self.module_name = o.fullname # Current module being processed
481485
self.path = o.path
@@ -646,12 +650,14 @@ def visit_func_def(self, o: FuncDef) -> None:
646650
if init_code:
647651
self.add(init_code)
648652

649-
if self._current_class is not None:
653+
if self._class_stack:
650654
if len(o.arguments):
651655
self_var = o.arguments[0].variable.name
652656
else:
653657
self_var = "self"
654-
class_info = ClassInfo(self._current_class.name, self_var)
658+
class_info: ClassInfo | None = None
659+
for class_def in self._class_stack:
660+
class_info = ClassInfo(class_def.name, self_var, parent=class_info)
655661
else:
656662
class_info = None
657663

@@ -741,7 +747,7 @@ def get_fullname(self, expr: Expression) -> str:
741747
return self.resolve_name(name)
742748

743749
def visit_class_def(self, o: ClassDef) -> None:
744-
self._current_class = o
750+
self._class_stack.append(o)
745751
self.method_names = find_method_names(o.defs.body)
746752
sep: int | None = None
747753
if self.is_top_level() and self._state != EMPTY:
@@ -786,8 +792,8 @@ def visit_class_def(self, o: ClassDef) -> None:
786792
self._state = CLASS
787793
self.method_names = set()
788794
self.processing_dataclass = False
795+
self._class_stack.pop(-1)
789796
self.processing_enum = False
790-
self._current_class = None
791797

792798
def get_base_types(self, cdef: ClassDef) -> list[str]:
793799
"""Get list of base classes for a class."""

mypy/stubgenc.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -787,7 +787,9 @@ def get_base_types(self, obj: type) -> list[str]:
787787
bases.append(base)
788788
return [self.strip_or_import(self.get_type_fullname(base)) for base in bases]
789789

790-
def generate_class_stub(self, class_name: str, cls: type, output: list[str]) -> None:
790+
def generate_class_stub(
791+
self, class_name: str, cls: type, output: list[str], parent_class: ClassInfo | None = None
792+
) -> None:
791793
"""Generate stub for a single class using runtime introspection.
792794
793795
The result lines will be appended to 'output'. If necessary, any
@@ -808,7 +810,9 @@ def generate_class_stub(self, class_name: str, cls: type, output: list[str]) ->
808810
self.record_name(class_name)
809811
self.indent()
810812

811-
class_info = ClassInfo(class_name, "", getattr(cls, "__doc__", None), cls)
813+
class_info = ClassInfo(
814+
class_name, "", getattr(cls, "__doc__", None), cls, parent=parent_class
815+
)
812816

813817
for attr, value in items:
814818
# use unevaluated descriptors when dealing with property inspection
@@ -843,7 +847,7 @@ def generate_class_stub(self, class_name: str, cls: type, output: list[str]) ->
843847
class_info,
844848
)
845849
elif inspect.isclass(value) and self.is_defined_in_module(value):
846-
self.generate_class_stub(attr, value, types)
850+
self.generate_class_stub(attr, value, types, parent_class=class_info)
847851
else:
848852
attrs.append((attr, value))
849853

mypy/stubutil.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -306,12 +306,18 @@ def args_str(self, args: Iterable[Type]) -> str:
306306

307307
class ClassInfo:
308308
def __init__(
309-
self, name: str, self_var: str, docstring: str | None = None, cls: type | None = None
309+
self,
310+
name: str,
311+
self_var: str,
312+
docstring: str | None = None,
313+
cls: type | None = None,
314+
parent: ClassInfo | None = None,
310315
) -> None:
311316
self.name = name
312317
self.self_var = self_var
313318
self.docstring = docstring
314319
self.cls = cls
320+
self.parent = parent
315321

316322

317323
class FunctionContext:
@@ -334,7 +340,13 @@ def __init__(
334340
def fullname(self) -> str:
335341
if self._fullname is None:
336342
if self.class_info:
337-
self._fullname = f"{self.module_name}.{self.class_info.name}.{self.name}"
343+
parents = []
344+
class_info: ClassInfo | None = self.class_info
345+
while class_info is not None:
346+
parents.append(class_info.name)
347+
class_info = class_info.parent
348+
namespace = ".".join(reversed(parents))
349+
self._fullname = f"{self.module_name}.{namespace}.{self.name}"
338350
else:
339351
self._fullname = f"{self.module_name}.{self.name}"
340352
return self._fullname

mypy/test/teststubgen.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
from mypy.stubgenc import InspectionStubGenerator, infer_c_method_args
3939
from mypy.stubutil import (
4040
ClassInfo,
41+
FunctionContext,
4142
common_dir_prefix,
4243
infer_method_ret_type,
4344
remove_misplaced_type_comments,
@@ -612,6 +613,16 @@ def test_common_dir_prefix_win(self) -> None:
612613
assert common_dir_prefix([r"foo\bar/x.pyi"]) == r"foo\bar"
613614
assert common_dir_prefix([r"foo/bar/x.pyi"]) == r"foo\bar"
614615

616+
def test_function_context_nested_classes(self):
617+
ctx = FunctionContext(
618+
module_name="spangle",
619+
name="foo",
620+
class_info=ClassInfo(
621+
name="Nested", self_var="self", parent=ClassInfo(name="Parent", self_var="self")
622+
),
623+
)
624+
assert ctx.fullname == "spangle.Parent.Nested.foo"
625+
615626

616627
class StubgenHelpersSuite(unittest.TestCase):
617628
def test_is_blacklisted_path(self) -> None:

0 commit comments

Comments
 (0)