Rework some of the analysis statement gathering.

jedi/api/__init__.py

@@ -535,11 +535,7 @@ class Script(object):
         for n in imp_names:
             imports.ImportWrapper(self._evaluator, n).follow()
         for node in sorted(nodes, key=lambda obj: obj.start_pos):
-            #if not (isinstance(stmt.parent, pr.ForFlow) and stmt.parent.set_stmt == stmt):
+            check_types(self._evaluator.eval_element(node))
-            if node.type == 'expr_stmt':
-                check_types(self._evaluator.eval_statement(node))
-            else:
-                check_types(self._evaluator.eval_element(node))
 
         for dec_func in decorated_funcs:
             er.Function(self._evaluator, dec_func).get_decorated_func()

jedi/evaluate/analysis.py

@@ -240,16 +240,33 @@ def get_module_statements(module):
                                     nodes.append(argument)
             return nodes
 
-    def add_stmts(stmts):
+    def add_nodes(nodes):
         new = set()
-        for stmt in stmts:
+        for node in nodes:
-            if stmt.type == 'expr_stmt':
+            if isinstance(node, pr.Flow):
-                new.add(stmt)
+                # Pick the suite/simple_stmt.
+                new |= add_nodes(node.children[-1].children)
+            elif node.type in ('simple_stmt', 'suite'):
+                new |= add_nodes(node.children)
+            elif node.type in ('return_stmt', 'yield_expr'):
+                try:
+                    new.add(node.children[1])
+                except IndexError:
+                    pass
+            elif node.type not in ('whitespace', 'operator', 'keyword',
+                                   'parameters', 'decorated') \
+                    and not isinstance(node, (pr.ClassOrFunc, pr.Import)):
+                new.add(node)
 
-            for node in stmt.children:
+                try:
-                new.update(check_children(node))
+                    children = node.children
-                if node.type != 'keyword' and stmt.type != 'expr_stmt':
+                except AttributeError:
-                    new.add(node)
+                    pass
+                else:
+                    for next_node in children:
+                        new.update(check_children(node))
+                        if next_node.type != 'keyword' and node.type != 'expr_stmt':
+                            new.add(node)
         return new
 
     nodes = set()
@@ -260,8 +277,11 @@ def get_module_statements(module):
             import_names |= set(imp.get_defined_names())
             if imp.is_nested():
                 import_names |= set(path[-1] for path in imp.paths())
-        nodes |= add_stmts(scope.statements)
+
-        nodes |= add_stmts(r for r in scope.returns if r is not None)
+        children = scope.children
+        if isinstance(scope, pr.ClassOrFunc):
+            children = children[2:]  # We don't want to include the class name.
+        nodes |= add_nodes(children)
 
         for flow in scope.flows:
             if flow.type == 'for_stmt':

jedi/evaluate/finder.py

@@ -149,7 +149,7 @@ class NameFinder(object):
                 continue
 
             if isinstance(self.name_str, pr.Name):
-                origin_scope = self.name_str.get_definition().parent
+                origin_scope = self.name_str.get_parent_until(pr.Scope, reverse=True)
             else:
                 origin_scope = None
             if isinstance(stmt.parent, compiled.CompiledObject):

test/test_parser/test_parser.py

@@ -21,8 +21,9 @@ def test_user_statement_on_import():
 
 class TestCallAndName():
     def get_call(self, source):
-        stmt = Parser(load_grammar(), u(source)).module.statements[0]
+        # Get the simple_stmt and then the first one.
-        return stmt.children[0]
+        simple_stmt = Parser(load_grammar(), u(source)).module.children[0]
+        return simple_stmt.children[0]
 
     def test_name_and_call_positions(self):
         name = self.get_call('name\nsomething_else')