From fcf85065319277bdad6e36ee9ab4b087c310d98c Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Sun, 22 Dec 2019 17:18:10 +0100
Subject: [PATCH] Add Script().get_context, fixes #253

---
 CHANGELOG.rst                 |   1 +
 jedi/api/__init__.py          |  24 ++++++++
 test/test_api/test_context.py | 113 ++++++++++++++++++++++++++++++++++
 3 files changed, 138 insertions(+)
 create mode 100644 test/test_api/test_context.py

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index d3582119..c6c95656 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -7,6 +7,7 @@ Changelog
 +++++++++++++++++++
 
 
+- **Add** ``Script.get_context`` to get information where you currently are.
 - Big **Script API Changes**:
     - The line and column parameters of ``jedi.Script`` are now deprecated
     - ``completions`` deprecated, use ``complete`` instead
diff --git a/jedi/api/__init__.py b/jedi/api/__init__.py
index 2fac065d..7dc00e69 100644
--- a/jedi/api/__init__.py
+++ b/jedi/api/__init__.py
@@ -418,6 +418,30 @@ class Script(object):
         return [classes.Signature(self._inference_state, signature, call_details)
                 for signature in definitions.get_signatures()]
 
+    @validate_line_column
+    def get_context(self, line=None, column=None):
+        leaf = self._module_node.get_leaf_for_position((line, column), include_prefixes=True)
+        if leaf.start_pos > (line, column) or leaf.type == 'endmarker':
+            previous_leaf = leaf.get_previous_leaf()
+            if previous_leaf is not None:
+                leaf = previous_leaf
+
+        module_context = self._get_module_context()
+        context = module_context.create_context(leaf)
+        while context.name is None:
+            context = context.parent_context  # comprehensions
+
+        definition = classes.Definition(self._inference_state, context.name)
+        while definition.type != 'module':
+            name = definition._name  # TODO private access
+            tree_name = name.tree_name
+            if tree_name is not None:  # Happens with lambdas.
+                scope = tree_name.get_definition()
+                if scope.start_pos[1] < column:
+                    break
+            definition = definition.parent()
+        return definition
+
     def _analysis(self):
         self._inference_state.is_analysis = True
         self._inference_state.analysis_modules = [self._module_node]
diff --git a/test/test_api/test_context.py b/test/test_api/test_context.py
new file mode 100644
index 00000000..82297c11
--- /dev/null
+++ b/test/test_api/test_context.py
@@ -0,0 +1,113 @@
+import pytest
+
+
+def _iter_hierarchy(context):
+    def iter(context):
+        while context is not None:
+            yield context
+            context = context.parent()
+
+    return reversed(list(iter(context)))
+
+
+func_code = '''\
+def func1(x, y):
+    pass
+
+def func2():
+ what ?
+i = 3
+    
+def func3():
+    1'''
+cls_code = '''\
+class Foo:
+    def x():
+        def y():
+            pass
+'''
+cls_nested = '''\
+class C:
+    class D:
+        def f():
+            pass
+'''
+lambda_ = '''\
+def x():
+    (lambda x:
+     lambda: y
+    )
+'''
+comprehension = '''
+def f(x):
+    [x
+    for
+    x
+    in x
+    ]'''
+
+with_brackets = '''\
+def x():
+    [
+
+    ]
+'''
+
+
+@pytest.mark.parametrize(
+    'code, line, column, full_name, expected_parents', [
+        ('', None, None, 'myfile', []),
+        (' ', None, 0, 'myfile', []),
+
+        (func_code, 1, 0, 'myfile', []),
+        (func_code, 1, None, 'myfile.func1', ['func1']),
+        #(func_code, 1, 4, 'myfile.func1', ['func1']),
+        #(func_code, 1, 10, 'myfile.func1', ['func1']),
+
+        (func_code, 3, 0, 'myfile', []),
+        (func_code, 5, None, 'myfile.func2', ['func2']),
+        (func_code, 6, None, 'myfile', []),
+        (func_code, 7, None, 'myfile', []),
+        (func_code, 9, None, 'myfile.func3', ['func3']),
+
+        (cls_code, None, None, 'myfile', []),
+        (cls_code + ' ', None, None, 'myfile.Foo', ['Foo']),
+        (cls_code + ' ' * 3, None, None, 'myfile.Foo', ['Foo']),
+        (cls_code + ' ' * 4, None, None, 'myfile.Foo', ['Foo']),
+        (cls_code + ' ' * 5, None, None, 'myfile.Foo.x', ['Foo', 'x']),
+        (cls_code + ' ' * 8, None, None, 'myfile.Foo.x', ['Foo', 'x']),
+        (cls_code + ' ' * 12, None, None, None, ['Foo', 'x', 'y']),
+
+        (cls_code, 4, 0, 'myfile', []),
+        (cls_code, 4, 3, 'myfile.Foo', ['Foo']),
+        (cls_code, 4, 4, 'myfile.Foo', ['Foo']),
+        (cls_code, 4, 5, 'myfile.Foo.x', ['Foo', 'x']),
+        (cls_code, 4, 8, 'myfile.Foo.x', ['Foo', 'x']),
+        (cls_code, 4, 12, None, ['Foo', 'x', 'y']),
+
+        (cls_nested, 4, None, 'myfile.C.D.f', ['C', 'D', 'f']),
+        (cls_nested, 4, 3, 'myfile.C', ['C']),
+
+        (lambda_, 2, 9, 'myfile.x', ['x']),  # the lambda keyword
+        (lambda_, 2, 13, 'myfile.x', ['x']),  # the lambda param
+        (lambda_, 3, 0, 'myfile', []),  # Within brackets, but they are ignored.
+        (lambda_, 3, 8, 'myfile.x', ['x']),
+        (lambda_, 3, None, 'myfile.x', ['x']),
+
+        (comprehension, 2, None, 'myfile.f', ['f']),
+        (comprehension, 3, None, 'myfile.f', ['f']),
+        (comprehension, 4, None, 'myfile.f', ['f']),
+        (comprehension, 5, None, 'myfile.f', ['f']),
+        (comprehension, 6, None, 'myfile.f', ['f']),
+
+        # Brackets are just ignored.
+        (with_brackets, 3, None, 'myfile', []),
+        (with_brackets, 4, 4, 'myfile.x', ['x']),
+        (with_brackets, 4, 5, 'myfile.x', ['x']),
+    ]
+)
+def test_context(Script, code, line, column, full_name, expected_parents):
+    context = Script(code, path='/foo/myfile.py').get_context(line, column)
+    assert context.full_name == full_name
+    parent_names = [d.name for d in _iter_hierarchy(context)]
+    assert parent_names == ['myfile'] + expected_parents