Skip to content

Commit

Permalink
Add support for try blocks in stubs
Browse files Browse the repository at this point in the history
On a local test I saw that kiwisolver/_cext.pyi and
rapidfuzz/process.pyi contain try blocks. Not standards-compliant
but I'd rather have typeshed-client not crash on this.
  • Loading branch information
JelleZijlstra committed Jul 12, 2024
1 parent 8f622f5 commit f9dcb32
Show file tree
Hide file tree
Showing 4 changed files with 27 additions and 0 deletions.
9 changes: 9 additions & 0 deletions tests/test.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,15 @@ def test_dot_import(self) -> None:
names["f"].ast, typeshed_client.ImportedName(path, "f")
)

def test_try(self) -> None:
ctx = get_context((3, 10))
names = get_stub_names("tryexcept", search_context=ctx)
assert names is not None
self.assertEqual(set(names), {"np", "f", "x"})
self.check_nameinfo(names, "np", typeshed_client.ImportedName)
self.check_nameinfo(names, "f", ast.FunctionDef)
self.check_nameinfo(names, "x", ast.AnnAssign)

def check_nameinfo(
self,
names: typeshed_client.NameDict,
Expand Down
1 change: 1 addition & 0 deletions tests/typeshed/VERSIONS
Original file line number Diff line number Diff line change
Expand Up @@ -17,3 +17,4 @@ about: 3.5-
importabout: 3.5-
tupleall: 3.5-
starimportall: 3.5-
tryexcept: 3.5-
9 changes: 9 additions & 0 deletions tests/typeshed/tryexcept.pyi
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
try:
import numpy as np

def f(x: np.int64) -> np.int64: ...

except ImportError:
pass
finally:
x: int
8 changes: 8 additions & 0 deletions typeshed_client/parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -252,6 +252,14 @@ def visit_If(self, node: ast.If) -> Iterable[NameInfo]:
for stmt in node.orelse:
yield from self.visit(stmt)

def visit_Try(self, node: ast.Try) -> Iterable[NameInfo]:
# try-except sometimes gets used with conditional imports. We assume
# the try block is always executed.
for stmt in node.body:
yield from self.visit(stmt)
for stmt in node.finalbody:
yield from self.visit(stmt)

def visit_Assert(self, node: ast.Assert) -> Iterable[NameInfo]:
visitor = _LiteralEvalVisitor(self.ctx)
value = visitor.visit(node.test)
Expand Down

0 comments on commit f9dcb32

Please sign in to comment.