Fix inference bug with If node

TheCavaco · Sep 22, 2022 · 823f0d1 · 823f0d1
1 parent d90cf6b
commit 823f0d1
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/py2many/inference.py b/py2many/inference.py
@@ -958,6 +958,12 @@ def _visit_val_branch(self, node: ast.If, ann, ann_node_id):
             ann_node.annotation = ast.Name(id="None")
             for n in node.orelse:
                 self.visit(n)
+        else:
+            node.test.annotation = ann
+            for n in node.body:
+                self.visit(n)
+            for n in node.orelse:
+                self.visit(n)
 
     def _visit_none_branch(self, node: ast.If, ann, ann_node_id):
         if self._is_optional(ann):
@@ -969,6 +975,12 @@ def _visit_none_branch(self, node: ast.If, ann, ann_node_id):
             ann_node.annotation = ann.slice
             for n in node.orelse:
                 self.visit(n)
+        else:
+            node.test.annotation = ann
+            for n in node.body:
+                self.visit(n)
+            for n in node.orelse:
+                self.visit(n)
 
     def _is_optional(self, annotation):
         is_optional = lambda x: get_id(x) == "Optional"