From e4170d65b7b7e39a313f6b5452e552419926d0fa Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Wed, 15 May 2019 21:55:54 +0200
Subject: [PATCH] Make namespace folders work with stubs

---
 jedi/evaluate/base_context.py                 |  3 +++
 jedi/evaluate/context/namespace.py            |  3 +++
 jedi/evaluate/gradual/typeshed.py             | 15 ++++++++++-----
 .../stub_folder/stub_only_folder/__init__.pyi |  1 +
 .../stub_folder/with_stub_folder/__init__.py  |  2 ++
 .../stub_folder/with_stub_folder/__init__.pyi |  2 ++
 test/completion/stubs.py                      | 19 ++++++++++++++++---
 7 files changed, 37 insertions(+), 8 deletions(-)
 create mode 100644 test/completion/stub_folder/stub_only_folder/__init__.pyi
 create mode 100644 test/completion/stub_folder/with_stub_folder/__init__.py
 create mode 100644 test/completion/stub_folder/with_stub_folder/__init__.pyi

diff --git a/jedi/evaluate/base_context.py b/jedi/evaluate/base_context.py
index 1a84c5d1..5327763d 100644
--- a/jedi/evaluate/base_context.py
+++ b/jedi/evaluate/base_context.py
@@ -158,6 +158,9 @@ class Context(HelperContextMixin, BaseContext):
     def is_module(self):
         return False
 
+    def is_namespace(self):
+        return False
+
     def is_compiled(self):
         return False
 
diff --git a/jedi/evaluate/context/namespace.py b/jedi/evaluate/context/namespace.py
index c973fb2f..12c8b3a6 100644
--- a/jedi/evaluate/context/namespace.py
+++ b/jedi/evaluate/context/namespace.py
@@ -54,6 +54,9 @@ class ImplicitNamespaceContext(Context, SubModuleDictMixin):
     def py__name__(self):
         return self._fullname
 
+    def is_namespace(self):
+        return True
+
     def is_stub(self):
         return False
 
diff --git a/jedi/evaluate/gradual/typeshed.py b/jedi/evaluate/gradual/typeshed.py
index 8b1acb55..d0cd34ed 100644
--- a/jedi/evaluate/gradual/typeshed.py
+++ b/jedi/evaluate/gradual/typeshed.py
@@ -160,12 +160,18 @@ def _try_to_load_stub(evaluator, import_names, actual_context_set,
             pass
         else:
             file_path = method()
-            if file_path is not None and file_path.endswith('.py'):
+            file_paths = []
+            if c.is_namespace():
+                file_paths = [os.path.join(p, '__init__.pyi') for p in c.py__path__()]
+            elif file_path is not None and file_path.endswith('.py'):
+                file_paths = [file_path + 'i']
+
+            for file_path in file_paths:
                 m = _try_to_load_stub_from_file(
                     evaluator,
                     actual_context_set,
                     # The file path should end with .pyi
-                    file_path + 'i',
+                    file_path,
                     import_names
                 )
                 if m is not None:
@@ -186,14 +192,13 @@ def _try_to_load_stub(evaluator, import_names, actual_context_set,
         else:
             check_path = sys_path
             names_for_path = import_names
-        names_for_path = names_for_path[:-1] + (names_for_path[-1] + '.pyi',)
 
         for p in check_path:
             m = _try_to_load_stub_from_file(
                 evaluator,
                 actual_context_set,
-                os.path.join(p, *names_for_path),
-                import_names
+                os.path.join(p, *names_for_path) + '.pyi',
+                import_names,
             )
             if m is not None:
                 return m
diff --git a/test/completion/stub_folder/stub_only_folder/__init__.pyi b/test/completion/stub_folder/stub_only_folder/__init__.pyi
new file mode 100644
index 00000000..77cf2466
--- /dev/null
+++ b/test/completion/stub_folder/stub_only_folder/__init__.pyi
@@ -0,0 +1 @@
+in_stub_only_folder: int
diff --git a/test/completion/stub_folder/with_stub_folder/__init__.py b/test/completion/stub_folder/with_stub_folder/__init__.py
new file mode 100644
index 00000000..4201289b
--- /dev/null
+++ b/test/completion/stub_folder/with_stub_folder/__init__.py
@@ -0,0 +1,2 @@
+in_with_stub_both_folder = 5
+in_with_stub_python_folder = 8
diff --git a/test/completion/stub_folder/with_stub_folder/__init__.pyi b/test/completion/stub_folder/with_stub_folder/__init__.pyi
new file mode 100644
index 00000000..ea7ec38c
--- /dev/null
+++ b/test/completion/stub_folder/with_stub_folder/__init__.pyi
@@ -0,0 +1,2 @@
+in_with_stub_both_folder: str
+in_with_stub_stub_folder: float
diff --git a/test/completion/stubs.py b/test/completion/stubs.py
index cccf0d81..437be1c1 100644
--- a/test/completion/stubs.py
+++ b/test/completion/stubs.py
@@ -1,11 +1,24 @@
-from stub_folder import with_stub, stub_only
+from stub_folder import with_stub, stub_only, with_stub_folder, stub_only_folder
+
+# Just files
 
 #? 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
+
+
+# Folders
+
+#? int()
+stub_only_folder.in_stub_only_folder
+#? str()
+with_stub_folder.in_with_stub_both_folder
+#?
+with_stub_folder.in_with_stub_python_folder
+#? float()
+with_stub_folder.in_with_stub_stub_folder