Skip to content

Commit

Permalink
[stubgen] Fix crash on literal class-level keywords (#17663)
Browse files Browse the repository at this point in the history
Closes #17661
  • Loading branch information
sobolevn authored Aug 12, 2024
1 parent f0c394a commit 5b46f1c
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 2 deletions.
22 changes: 20 additions & 2 deletions mypy/stubgen.py
Original file line number Diff line number Diff line change
Expand Up @@ -305,9 +305,26 @@ def visit_name_expr(self, node: NameExpr) -> str:
def visit_member_expr(self, o: MemberExpr) -> str:
return self._visit_ref_expr(o)

def visit_str_expr(self, node: StrExpr) -> str:
def _visit_literal_node(
self, node: StrExpr | BytesExpr | IntExpr | FloatExpr | ComplexExpr
) -> str:
return repr(node.value)

def visit_str_expr(self, node: StrExpr) -> str:
return self._visit_literal_node(node)

def visit_bytes_expr(self, node: BytesExpr) -> str:
return f"b{self._visit_literal_node(node)}"

def visit_int_expr(self, node: IntExpr) -> str:
return self._visit_literal_node(node)

def visit_float_expr(self, node: FloatExpr) -> str:
return self._visit_literal_node(node)

def visit_complex_expr(self, node: ComplexExpr) -> str:
return self._visit_literal_node(node)

def visit_index_expr(self, node: IndexExpr) -> str:
base_fullname = self.stubgen.get_fullname(node.base)
if base_fullname == "typing.Union":
Expand Down Expand Up @@ -805,7 +822,8 @@ def get_base_types(self, cdef: ClassDef) -> list[str]:
for name, value in cdef.keywords.items():
if name == "metaclass":
continue # handled separately
base_types.append(f"{name}={value.accept(p)}")
processed_value = value.accept(p) or "..." # at least, don't crash
base_types.append(f"{name}={processed_value}")
return base_types

def get_class_decorators(self, cdef: ClassDef) -> list[str]:
Expand Down
10 changes: 10 additions & 0 deletions test-data/unit/stubgen.test
Original file line number Diff line number Diff line change
Expand Up @@ -4405,3 +4405,13 @@ Y = int | None
Z = Incomplete
W = int | str | None
R = type[int | str] | None

[case testClassInheritanceWithKeywordsConstants]
class Test(Whatever, a=1, b='b', c=True, d=1.5, e=None, f=1j, g=b'123'): ...
[out]
class Test(Whatever, a=1, b='b', c=True, d=1.5, e=None, f=1j, g=b'123'): ...

[case testClassInheritanceWithKeywordsDynamic]
class Test(Whatever, keyword=SomeName * 2, attr=SomeName.attr): ...
[out]
class Test(Whatever, keyword=SomeName * 2, attr=SomeName.attr): ...

0 comments on commit 5b46f1c

Please sign in to comment.