added import completion

Trundle · Trundle · commit 7f9a9b258d7b · 2009-05-30T12:21:48.000+02:00
diff --git a/bpython/cli.py b/bpython/cli.py
@@ -55,6 +55,9 @@
 from pygments.lexers import PythonLexer
 from bpython.formatter import BPythonFormatter
 
+# This for import completion
+from bpython import importcompletion
+
 # And these are used for argspec.
 from pyparsing import Forward, Suppress, QuotedString, dblQuotedString, \
     Group, OneOrMore, ZeroOrMore, Literal, Optional, Word, \
@@ -535,26 +538,31 @@ def _complete(self, unused_tab=False):
         if not cw:
             self.matches = []
 
-        try:
-            self.completer.complete(cw, 0)
-        except Exception:
+        # Check for import completion
+        e = False
+        matches = importcompletion.complete(self.s, cw)
+        if not matches:
+            # Nope, no import, continue with normal completion
+            try:
+                self.completer.complete(cw, 0)
+            except Exception:
 # This sucks, but it's either that or list all the exceptions that could
 # possibly be raised here, so if anyone wants to do that, feel free to send me
 # a patch. XXX: Make sure you raise here if you're debugging the completion
 # stuff !
-            e = True
-        else:
-            e = False
+                e = True
+            else:
+                matches = self.completer.matches
 
-        if e or not self.completer.matches:
+        if e or not matches:
             self.matches = []
             if not self.argspec:
                 self.scr.redrawwin()
                 return False
 
-        if not e and self.completer.matches:
+        if not e and matches:
 # remove duplicates and restore order
-            self.matches = sorted(set(self.completer.matches))
+            self.matches = sorted(set(matches))
 
         if len(self.matches) == 1 and not OPTS.auto_display_list:
             self.list_win_visible = True
diff --git a/bpython/importcompletion.py b/bpython/importcompletion.py
@@ -0,0 +1,106 @@
+# The MIT License
+#
+# Copyright (c) 2009 Andreas Stuehrk
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+
+
+import imp
+import os
+import sys
+
+
+# The cached list of all known modules
+modules = list()
+
+
+def complete(line, cw):
+    """Construct a full list of possibly completions for imports."""
+    tokens = line.split()
+    if tokens[0] in ['from', 'import']:
+        if tokens[0] == 'from':
+            completing_from = True
+            if len(tokens) > 3:
+                if '.' in cw:
+                    # This will result in a SyntaxError, so do not return
+                    # any matches
+                    return list()
+                cw = '%s.%s' % (tokens[1], cw)
+            elif len(tokens) == 3:
+                return ['import']
+        else:
+            completing_from = False
+        
+        matches = list()
+        for name in modules:
+            if not (name.startswith(cw) and name.find('.', len(cw)) == -1):
+                continue
+            if completing_from:
+                name = name[len(tokens[1]) + 1:]
+            matches.append(name)
+        return matches
+    else:
+        return list()
+
+
+def find_modules(path):
+    """Find all modules (and packages) for a given directory."""
+    if not os.path.isdir(path):
+        # Perhaps a zip file
+        return
+
+    for name in os.listdir(path):
+        if not any(name.endswith(suffix[0]) for suffix in imp.get_suffixes()):
+            # Possibly a package
+            if '.' in name:
+                continue
+        name = os.path.splitext(name)[0]
+        try:
+            fo, pathname, _ = imp.find_module(name, [path])
+        except ImportError:
+            continue
+        else:
+            if fo is not None:
+                fo.close()
+            else:
+                # Yay, package
+                for subname in find_modules(pathname):
+                    if subname != '__init__':
+                        yield '%s.%s' % (name, subname)
+            yield name
+
+
+def find_all_modules(path=None):
+    """Return a list with all modules in `path`, which should be a list of
+    directory names. If path is not given, sys.path will be used."""
+    modules = set()
+    if path is None:
+        modules.update(sys.builtin_module_names)
+        path = sys.path
+
+    for p in path:
+        modules.update(find_modules(p))
+    
+    return modules
+
+
+def reload():
+    """Refresh the list of known modules."""
+    global modules
+    modules = find_all_modules()
