From f9dcb32be892f6a41e9569dfbf44c0c8167cef6d Mon Sep 17 00:00:00 2001
From: Jelle Zijlstra <jelle.zijlstra@gmail.com>
Date: Fri, 12 Jul 2024 15:33:03 -0700
Subject: [PATCH] Add support for try blocks in stubs

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.
---
 tests/test.py                | 9 +++++++++
 tests/typeshed/VERSIONS      | 1 +
 tests/typeshed/tryexcept.pyi | 9 +++++++++
 typeshed_client/parser.py    | 8 ++++++++
 4 files changed, 27 insertions(+)
 create mode 100644 tests/typeshed/tryexcept.pyi

diff --git a/tests/test.py b/tests/test.py
index 1b4a5dbf..9ca80997 100644
--- a/tests/test.py
+++ b/tests/test.py
@@ -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,
diff --git a/tests/typeshed/VERSIONS b/tests/typeshed/VERSIONS
index 766159de..5fb69683 100644
--- a/tests/typeshed/VERSIONS
+++ b/tests/typeshed/VERSIONS
@@ -17,3 +17,4 @@ about: 3.5-
 importabout: 3.5-
 tupleall: 3.5-
 starimportall: 3.5-
+tryexcept: 3.5-
diff --git a/tests/typeshed/tryexcept.pyi b/tests/typeshed/tryexcept.pyi
new file mode 100644
index 00000000..a96b28f7
--- /dev/null
+++ b/tests/typeshed/tryexcept.pyi
@@ -0,0 +1,9 @@
+try:
+    import numpy as np
+
+    def f(x: np.int64) -> np.int64: ...
+
+except ImportError:
+    pass
+finally:
+    x: int
diff --git a/typeshed_client/parser.py b/typeshed_client/parser.py
index a2a83e58..97393d54 100644
--- a/typeshed_client/parser.py
+++ b/typeshed_client/parser.py
@@ -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)