From 5ff3e4d1d1794c80801f9a4093857bf5adb68aab Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Mon, 13 May 2019 10:13:59 +0200
Subject: [PATCH] Implement stub tests and a first iteration of loading them
 from some random place

---
 jedi/evaluate/gradual/typeshed.py         | 55 +++++++++++++++++------
 test/completion/stub_folder/stub_only.pyi |  1 +
 test/completion/stub_folder/with_stub.py  |  2 +
 test/completion/stub_folder/with_stub.pyi |  2 +
 test/completion/stubs.py                  | 11 +++++
 5 files changed, 57 insertions(+), 14 deletions(-)
 create mode 100644 test/completion/stub_folder/stub_only.pyi
 create mode 100644 test/completion/stub_folder/with_stub.py
 create mode 100644 test/completion/stub_folder/with_stub.pyi
 create mode 100644 test/completion/stubs.py

diff --git a/jedi/evaluate/gradual/typeshed.py b/jedi/evaluate/gradual/typeshed.py
index d20b11ef..2a4ce755 100644
--- a/jedi/evaluate/gradual/typeshed.py
+++ b/jedi/evaluate/gradual/typeshed.py
@@ -128,6 +128,35 @@ def _try_to_load_stub(evaluator, actual_context_set, parent_module_context, impo
     except KeyError:
         pass
 
+    for c in actual_context_set:
+        try:
+            method = c.py__file__
+        except AttributeError:
+            pass
+        else:
+            file_path = method()
+            if file_path is not None and file_path.endswith('.py'):
+                m = _try_to_load_stub_from_file(
+                    evaluator,
+                    actual_context_set,
+                    # The file path should end with .pyi
+                    file_path + 'i',
+                    import_names
+                )
+                if m is not None:
+                    return m
+
+    m = _load_from_typeshed(evaluator, actual_context_set, parent_module_context, import_names)
+    if m is not None:
+        return m
+
+    evaluator.stub_module_cache[import_names] = None
+    # If no stub is found, just return the default.
+
+    return None
+
+
+def _load_from_typeshed(evaluator, actual_context_set, parent_module_context, import_names):
     import_name = import_names[-1]
     map_ = None
     if len(import_names) == 1:
@@ -143,22 +172,20 @@ def _try_to_load_stub(evaluator, actual_context_set, parent_module_context, impo
     if map_ is not None:
         path = map_.get(import_name)
         if path is not None:
-            try:
-                stub_module_node = _load_stub(evaluator, path)
-            except FileNotFoundError:
-                # The file has since been removed after looking for it.
-                # TODO maybe empty cache?
-                pass
-            else:
-                return create_stub_module(
-                    evaluator, actual_context_set, stub_module_node, path,
-                    import_names
-                )
-    evaluator.stub_module_cache[import_names] = None
-    # If no stub is found, just return the default.
+            return _try_to_load_stub_from_file(evaluator, actual_context_set, path, import_names)
 
-    return None
 
+def _try_to_load_stub_from_file(evaluator, actual_context_set, path, import_names):
+    try:
+        stub_module_node = _load_stub(evaluator, path)
+    except FileNotFoundError:
+        # The file that you're looking for doesn't exist (anymore).
+        return None
+    else:
+        return create_stub_module(
+            evaluator, actual_context_set, stub_module_node, path,
+            import_names
+        )
 
 def create_stub_module(evaluator, actual_context_set, stub_module_node, path, import_names):
     if import_names == ('typing',):
diff --git a/test/completion/stub_folder/stub_only.pyi b/test/completion/stub_folder/stub_only.pyi
new file mode 100644
index 00000000..4f2f4239
--- /dev/null
+++ b/test/completion/stub_folder/stub_only.pyi
@@ -0,0 +1 @@
+in_stub_only: int
diff --git a/test/completion/stub_folder/with_stub.py b/test/completion/stub_folder/with_stub.py
new file mode 100644
index 00000000..7f064dca
--- /dev/null
+++ b/test/completion/stub_folder/with_stub.py
@@ -0,0 +1,2 @@
+in_with_stub_both = 5
+in_with_stub_python = 8
diff --git a/test/completion/stub_folder/with_stub.pyi b/test/completion/stub_folder/with_stub.pyi
new file mode 100644
index 00000000..7a3f3ecf
--- /dev/null
+++ b/test/completion/stub_folder/with_stub.pyi
@@ -0,0 +1,2 @@
+in_with_stub_both: str
+in_with_stub_stub: float
diff --git a/test/completion/stubs.py b/test/completion/stubs.py
new file mode 100644
index 00000000..cccf0d81
--- /dev/null
+++ b/test/completion/stubs.py
@@ -0,0 +1,11 @@
+from stub_folder import with_stub, stub_only
+
+#? int()
+stub_only.in_stub_only
+
+#? str()
+with_stub.in_with_stub_both
+#? int()
+with_stub.in_with_stub_python
+#? float()
+with_stub.in_with_stub_stub