Repository: kislyuk/argcomplete

Issue:
Support for python -m module
Is it possible for argcomplete to work for modules run as `python -m module`? This is necessary to do once you start organizing your code in a module (rather than just a single script) to make relative imports work correctly, but I don't really want to make a separate entry point script if I don't have to. 

I am using global completion and it isn't working, at least not out of the box. 

Comments:

Patch:
diff --git a/README.rst b/README.rst
index 1620461..b77e533 100644
--- a/README.rst
+++ b/README.rst
@@ -212,6 +212,10 @@ In global completion mode, you don't have to register each argcomplete-capable e
 will look for the string **PYTHON_ARGCOMPLETE_OK** in the first 1024 bytes of any executable that it's running
 completion for, and if it's found, follow the rest of the argcomplete protocol as described above.
 
+Additionally, completion is activated for scripts run as ``python <script>`` and ``python -m <module>``.
+This also works for alternate Python versions (e.g. ``python3`` and ``pypy``), as long as that version of Python has
+argcomplete installed.
+
 .. admonition:: Bash version compatibility
 
  Global completion requires bash support for ``complete -D``, which was introduced in bash 4.2. On OS X or older Linux
diff --git a/argcomplete/__init__.py b/argcomplete/__init__.py
index 029c6b8..25d4328 100644
--- a/argcomplete/__init__.py
+++ b/argcomplete/__init__.py
@@ -203,9 +203,13 @@ class CompletionFinder(object):
         comp_line = ensure_str(comp_line)
         cword_prequote, cword_prefix, cword_suffix, comp_words, last_wordbreak_pos = split_line(comp_line, comp_point)
 
-        if os.environ["_ARGCOMPLETE"] == "2":
-            # Shell hook recognized the first word as the interpreter; discard it
-            comp_words.pop(0)
+        # _ARGCOMPLETE is set by the shell script to tell us where comp_words
+        # should start, based on what we're completing.
+        # 1: <script> [args]
+        # 2: python <script> [args]
+        # 3: python -m <module> [args]
+        start = int(os.environ["_ARGCOMPLETE"]) - 1
+        comp_words = comp_words[start:]
 
         debug("\nLINE: '{l}'\nPREQUOTE: '{pq}'\nPREFIX: '{p}'".format(l=comp_line, pq=cword_prequote, p=cword_prefix),
               "\nSUFFIX: '{s}'".format(s=cword_suffix),
diff --git a/argcomplete/_check_module.py b/argcomplete/_check_module.py
new file mode 100644
index 0000000..9d14b3d
--- /dev/null
+++ b/argcomplete/_check_module.py
@@ -0,0 +1,49 @@
+import os
+import sys
+
+try:
+    from importlib.util import find_spec
+except ImportError:
+    from collections import namedtuple
+    from imp import find_module
+
+    ModuleSpec = namedtuple(
+        'ModuleSpec', ['origin', 'has_location', 'submodule_search_locations'])
+
+    def find_spec(name):
+        """Minimal implementation as required by `find`."""
+        f, path, _ = find_module(name)
+        has_location = path is not None
+        if f is None:
+            return ModuleSpec(None, has_location, [path])
+        f.close()
+        return ModuleSpec(path, has_location, None)
+
+
+def find(name):
+    names = name.split('.')
+    spec = find_spec(names[0])
+    if not spec.has_location:
+        raise Exception('cannot locate file')
+    if spec.submodule_search_locations is None:
+        if len(names) != 1:
+            raise Exception('{} is not a package'.format(names[0]))
+        return spec.origin
+    if len(spec.submodule_search_locations) != 1:
+        raise Exception('expecting one search location')
+    path = os.path.join(spec.submodule_search_locations[0], *names[1:])
+    if os.path.isdir(path):
+        return os.path.join(path, '__main__.py')
+    else:
+        return path + '.py'
+
+
+def main():
+    with open(find(sys.argv[1])) as f:
+        head = f.read(1024)
+    if 'PYTHON_ARGCOMPLETE_OK' not in head:
+        raise Exception('marker not found')
+
+
+if __name__ == '__main__':
+    main()
diff --git a/argcomplete/bash_completion.d/python-argcomplete.sh b/argcomplete/bash_completion.d/python-argcomplete.sh
index 41bd53b..3e520a5 100644
--- a/argcomplete/bash_completion.d/python-argcomplete.sh
+++ b/argcomplete/bash_completion.d/python-argcomplete.sh
@@ -18,7 +18,13 @@ _python_argcomplete_global() {
 
     local ARGCOMPLETE=0
     if [[ "$executable" == python* ]] || [[ "$executable" == pypy* ]]; then
-        if [[ -f "${COMP_WORDS[1]}" ]] && (head -c 1024 "${COMP_WORDS[1]}" | grep --quiet "PYTHON_ARGCOMPLETE_OK") >/dev/null 2>&1; then
+        if [[ "${COMP_WORDS[1]}" == -m ]]; then
+            if "$executable" -m argcomplete._check_module "${COMP_WORDS[2]}" >/dev/null 2>&1; then
+                ARGCOMPLETE=3
+            else
+                return
+            fi
+        elif [[ -f "${COMP_WORDS[1]}" ]] && (head -c 1024 "${COMP_WORDS[1]}" | grep --quiet "PYTHON_ARGCOMPLETE_OK") >/dev/null 2>&1; then
             local ARGCOMPLETE=2
         else
             return
@@ -36,7 +42,7 @@ _python_argcomplete_global() {
         fi
     fi
 
-    if [[ $ARGCOMPLETE == 1 ]] || [[ $ARGCOMPLETE == 2 ]]; then
+    if [[ $ARGCOMPLETE != 0 ]]; then
         local IFS=$(echo -e '\v')
         COMPREPLY=( $(_ARGCOMPLETE_IFS="$IFS" \
             COMP_LINE="$COMP_LINE" \


Test patch:
diff --git a/test/test.py b/test/test.py
index ed86324..7a48798 100755
--- a/test/test.py
+++ b/test/test.py
@@ -17,6 +17,7 @@ from argcomplete import (
     CompletionFinder,
     split_line,
     ExclusiveCompletionFinder,
+    _check_module
 )
 from argcomplete.completers import FilesCompleter, DirectoriesCompleter
 from argcomplete.compat import USING_PYTHON2, str, sys_encoding, ensure_str, ensure_bytes
@@ -61,7 +62,7 @@ class TestArgcomplete(unittest.TestCase):
     def setUp(self):
         self._os_environ = os.environ
         os.environ = os.environ.copy()
-        os.environ["_ARGCOMPLETE"] = "yes"
+        os.environ["_ARGCOMPLETE"] = "1"
         os.environ["_ARC_DEBUG"] = "yes"
         os.environ["IFS"] = IFS
         os.environ["_ARGCOMPLETE_COMP_WORDBREAKS"] = COMP_WORDBREAKS
@@ -872,6 +873,64 @@ class TestSplitLine(unittest.TestCase):
         self.assertEqual(self.wordbreak('"b:c=d" '), None)
 
 
+class TestCheckModule(unittest.TestCase):
+    def setUp(self):
+        self.dir = TempDir(prefix="test_dir_module", dir=".")
+        self.dir.__enter__()
+        # There is some odd bug that seems to only come up in Python 3.4 where
+        # using "." in sys.path sometimes won't find modules, so we'll use the
+        # full path each time.
+        sys.path.insert(0, os.getcwd())
+
+    def tearDown(self):
+        sys.path.pop(0)
+        self.dir.__exit__()
+
+    def test_module(self):
+        self._mkfile('module.py')
+        path = _check_module.find('module')
+        self.assertEqual(path, os.path.abspath('module.py'))
+        self.assertNotIn('module', sys.modules)
+
+    def test_package(self):
+        os.mkdir('package')
+        self._mkfile('package/__init__.py')
+        self._mkfile('package/module.py')
+        path = _check_module.find('package.module')
+        self.assertEqual(path, os.path.abspath('package/module.py'))
+        self.assertNotIn('package', sys.modules)
+        self.assertNotIn('package.module', sys.modules)
+
+    def test_subpackage(self):
+        os.mkdir('package')
+        self._mkfile('package/__init__.py')
+        os.mkdir('package/subpackage')
+        self._mkfile('package/subpackage/__init__.py')
+        self._mkfile('package/subpackage/module.py')
+        path = _check_module.find('package.subpackage.module')
+        self.assertEqual(path, os.path.abspath('package/subpackage/module.py'))
+        self.assertNotIn('package', sys.modules)
+        self.assertNotIn('package.subpackage', sys.modules)
+        self.assertNotIn('package.subpackage.module', sys.modules)
+
+    def test_package_main(self):
+        os.mkdir('package')
+        self._mkfile('package/__init__.py')
+        self._mkfile('package/__main__.py')
+        path = _check_module.find('package')
+        self.assertEqual(path, os.path.abspath('package/__main__.py'))
+        self.assertNotIn('package', sys.modules)
+
+    def test_not_package(self):
+        self._mkfile('module.py')
+        with self.assertRaisesRegexp(Exception, 'module is not a package'):
+            _check_module.find('module.bad')
+        self.assertNotIn('module', sys.modules)
+
+    def _mkfile(self, path):
+        open(path, 'w').close()
+
+
 class _TestSh(object):
     """
     Contains tests which should work in any shell using argcomplete.
@@ -1040,6 +1099,16 @@ class TestBashGlobal(TestBash):
             self.assertIn('Permission denied', self.sh.run_command('./prog'))
             self.assertEqual(self.sh.run_command('python ./prog basic f\t'), 'foo\r\n')
 
+    def test_python_module(self):
+        """Test completing a module run with python -m."""
+        prog = os.path.join(TEST_DIR, 'prog')
+        with TempDir(prefix='test_dir_py', dir='.'):
+            os.mkdir('package')
+            open('package/__init__.py', 'w').close()
+            shutil.copy(prog, 'package/prog.py')
+            self.sh.run_command('cd ' + os.getcwd())
+            self.assertEqual(self.sh.run_command('python -m package.prog basic f\t'), 'foo\r\n')
+
 
 class TestTcsh(_TestSh, unittest.TestCase):
     expected_failures = [