More work on attribut completion on tab

etiennerichart · etiennerichart · commit d3fa4c9bbacc · 2020-10-19T14:11:03.000-05:00
diff --git a/bpython/autocomplete.py b/bpython/autocomplete.py
@@ -240,6 +240,9 @@ def locate(self, current_offset, line):
     def format(self, word):
         return after_last_dot(word)
 
+    def try_to_complete(self, module):
+        return importcompletion.try_to_import(module)
+
 
 class FilenameCompletion(BaseCompletionType):
     def __init__(self, mode=SIMPLE):
diff --git a/bpython/curtsiesfrontend/repl.py b/bpython/curtsiesfrontend/repl.py
@@ -37,6 +37,7 @@
 from bpython import autocomplete
 from bpython.translations import _
 from bpython.pager import get_pager_command
+from bpython.line import from_import_tab
 
 from bpython.curtsiesfrontend import replpainter as paint
 from bpython.curtsiesfrontend import sitefix
@@ -887,6 +888,15 @@ def only_whitespace_left_of_cursor():
             for unused in range(to_add):
                 self.add_normal_character(" ")
             return
+        
+        # if line is `from a import` and tab then import module and attr
+        line = self.current_line.split()
+        module_name = ""
+        if "from" in line and "import" in line:
+            module_name = from_import_tab(self.current_line)
+        if module_name:
+            self.status_bar.message("Importing module %s" % module_name)
+            self.completers[1].try_to_complete(module_name)
 
         # run complete() if we don't already have matches
         if len(self.matches_iter.matches) == 0:
diff --git a/bpython/importcompletion.py b/bpython/importcompletion.py
@@ -92,6 +92,19 @@ def module_attr_matches(name):
     return attr_matches(name, prefix="", only_modules=True)
 
 
+def try_to_import(module_name):
+    """If this hasn't been imported our goal is to add it to the set of modules and actually import it"""
+    if module_name not in sys.modules:
+        try:
+            module = __import__(module_name)
+            module_path = os.path.dirname(module.__file__)
+            for mod in find_modules(module_path):
+                modules.add("%s.%s" % (module_name, mod))
+            modules.add(module_name)
+        except ImportError:
+            pass
+
+
 def complete(cursor_offset, line):
     """Construct a full list of possibly completions for imports."""
     tokens = line.split()
@@ -205,7 +218,7 @@ def find_all_modules(path=None):
             p = os.curdir
         for module in find_modules(p):
             module = try_decode(module, "ascii")
-            if module is None or "numpy" in module:
+            if module is None:
                 continue
             modules.add(module)
             yield
diff --git a/bpython/line.py b/bpython/line.py
@@ -117,6 +117,19 @@ def current_object_attribute(cursor_offset, line):
     return None
 
 
+from_import_tab_re = LazyReCompile(r"from ([\w0-9_.]*)\s+import")
+
+
+def from_import_tab(line):
+    matches = from_import_tab_re.finditer(line)
+    try:
+        module = matches.__next__()
+    except StopIteration:
+        pass
+    if module:
+        return str(module[1])
+
+
 current_from_import_from_re = LazyReCompile(
     r"from ([\w0-9_.]*)(?:\s+import\s+([\w0-9_]+[,]?\s*)+)*"
 )
