Repository: reynoldsnlp/pipster (language: python)

Issue:
More robust way to get `# install` comments
`_experimental.autoinstall._get_deps()` needs a more robust way to get `# pip install` comments, perhaps a la https://stackoverflow.com/a/36055400/2903532


Discussion / hints:

Fix patch:
diff --git a/src/pip_inside/__init__.py b/src/pip_inside/__init__.py
index e72c929..2b9a59a 100644
--- a/src/pip_inside/__init__.py
+++ b/src/pip_inside/__init__.py
@@ -1,151 +1,3 @@
-from glob import glob
-import os
-from pathlib import Path
-import re
-import shlex
-import subprocess
-import sys
-from warnings import warn
-
-try:
-    from pip._internal.commands.install import InstallCommand
-    from pip._vendor.distlib.database import DistributionPath
-    install_cmd = InstallCommand(name='dummy', summary='Provides parse_args.')
-except ModuleNotFoundError:
-    raise ModuleNotFoundError('Please install pip for the current '
-                              'interpreter: (%s).' % sys.executable)
-
 from ._version import version as __version__  # noqa: F401
-
-
-pipfiles = []
-for i in range(4):  # TODO: 4 is completely arbitrary here. improve?
-    pipfiles.extend(glob(('..' + os.sep) * i + 'Pipfile'))
-if pipfiles:
-    pipfiles = [os.path.abspath(f) for f in pipfiles]
-    msg = ('Warning: the following Pipfiles will be bypassed by '
-           'pip_inside.install:\n\t' + '\n\t'.join(pipfiles))
-    warn(msg, stacklevel=2)
-
-
-def _parse_target(target):
-    """Parse install target down to the distribution name."""
-    # TODO: make more robust
-    if target.endswith('.whl'):
-        return Path(target).name.split('-')[0]
-    else:
-        # https://pip.pypa.io/en/stable/reference/requirement-specifiers/
-        return re.search(r'(.+?)(?:[!~<>=]{1,2}.+)?$', target).group(1)
-
-
-def install(*args, **kwargs):
-    """Install packages into the current environment.
-
-    Equivalent examples of command-line pip and pip_inside are grouped below.
-
-    METHOD 1: Single argument is exactly the same as command line interface,
-    beginning with 'pip install ...'
-
-    $ pip install --user --upgrade some_pkg
-    >>> install('pip install --user --upgrade some_pkg')
-
-    METHOD 2: Arguments from command-line implementation split on spaces
-
-    $ pip install some_pkg
-    >>> install('some_pkg')
-
-    $ pip install --user --upgrade some_pkg
-    >>> install('--user', '--upgrade', 'some_pkg')
-    >>> install(*'--user --upgrade some_pkg'.split())
-
-    If preferred, keyword-value arguments can also be used:
-
-    $ pip install -r requirements.txt
-    >>> install(r='requirements.txt')
-    >>> install('-r', 'requirements.txt')
-    >>> install(*'-r requirements.txt'.split())
-
-    $ pip install --no-index --find-links /local/dir/ some_pkg
-    # Note the use of '_' in the following keyword example.
-    >>> install('--no-index', 'some_pkg', find_links='/local/dir/')
-    >>> install('--no-index', '--find-links', '/local/dir/', 'some_pkg')
-    >>> install(*'--no-index --find-links /local/dir/ some_pkg'.split())
-
-    """
-    cli_args = _build_install_cmd(*args, **kwargs)
-    # use pip internals to isolate package names
-    _, req_targets = install_cmd.parse_args(cli_args)
-    assert req_targets[:2] == ['pip', 'install']
-    req_targets = [_parse_target(t) for t in req_targets[2:]]
-    dist_path = DistributionPath(include_egg=True)
-    target_providers = [d.modules
-                        for t in req_targets
-                        for d in dist_path.get_distributions()  # installed
-                        if t == d.name]
-    target_providers = [y for x in target_providers for y in x]  # flatten
-    target_origins = {}
-    for t in target_providers:
-        if t in sys.modules:
-            target_origins[sys.modules[t].__spec__.origin] = t
-    dist_path = DistributionPath(include_egg=True)
-    dists = [dist_path.get_distribution(t) for t in req_targets]
-    paths = set()
-    for dist in dists:
-        if dist is not None:
-            for path, _, _ in dist.list_installed_files():
-                paths.add(os.path.join(os.path.dirname(dist.path), path))
-    already_loaded = paths.intersection(target_origins)
-    already_loaded = [target_origins[origin] for origin in already_loaded]
-    print('Trying  ', ' '.join(cli_args), '  ...', file=sys.stderr)
-    cli_cmd = [sys.executable, "-m"] + cli_args
-    result = subprocess.run(cli_cmd, check=True)
-    if result.returncode == 0 and already_loaded:
-        warn('WARNING! The following modules were already loaded. Restart '
-             f'python to see changes:  {", ".join(already_loaded)}',
-             UserWarning)
-    return result
-
-
-def _build_install_cmd(*args, **kwargs):
-    if len(args) == 1 and args[0].startswith('pip install '):
-        return shlex.split(args[0])
-    else:
-        cli_args = ['pip', 'install']
-        # Keyword arguments are translated to CLI options
-        for raw_k, v in kwargs.items():
-            k = raw_k.replace('_', '-')  # Python identifiers -> CLI long names
-            append_value = isinstance(v, str)
-            if append_value:
-                # When arg value is str, append both it and option to CLI args
-                append_option = True
-                if not v:
-                    raise ValueError("Empty string passed as value for option "
-                                     "{}".format(k))
-            else:
-                # assume the value indicates whether to include a boolean flag.
-                # None->omit, true->include, false->include negated
-                append_option = v is not None
-                if k.startswith("no-"):
-                    # suggest `some-option=True` instead of
-                    # `no-some-option=False`
-                    raw_suffix = raw_k[3:]
-                    msg_template = ("Rather than '{}={!r}', "
-                                    "try '{{}}={{!r}}'".format(raw_k, v))
-                    if append_option:
-                        suggestion = msg_template.format(raw_suffix, not v)
-                    else:
-                        suggestion = msg_template.format(raw_suffix, None)
-                    raise ValueError(suggestion)
-                if append_option and not v:
-                    k = "no-" + k
-            if append_option:
-                if len(k) == 1:  # short flag
-                    option = "-" + k
-                else:  # long flag
-                    option = "--" + k
-                cli_args.append(option)
-                if append_value:
-                    cli_args.append(v)
-        # Positional arguments are passed directly as CLI arguments
-        cli_args += args
-    return cli_args
+from .autoinstall import *
+from .install import *
diff --git a/src/pip_inside/autoinstall.py b/src/pip_inside/autoinstall.py
new file mode 100644
index 0000000..21d0262
--- /dev/null
+++ b/src/pip_inside/autoinstall.py
@@ -0,0 +1,100 @@
+import ast
+from collections import namedtuple
+import distutils.sysconfig as sysconfig
+import inspect
+from modulefinder import ModuleFinder
+import os
+import re
+import sys
+
+from . import install
+__all__ = ['autoinstall']
+
+
+def get_stdlib_module_names():
+    if sys.version_info >= (3, 10):
+        return sys.stdlib_module_names
+    else:
+        stdlib_module_names = set()
+        std_lib = sysconfig.get_python_lib(standard_lib=True)
+
+        for top, dirs, files in os.walk(std_lib):
+            for nm in files:
+                prefix = top[len(std_lib)+1:]
+                if prefix[:13] == 'site-packages':
+                    continue
+                if nm == '__init__.py':
+                    stdlib_module_names.add(top[len(std_lib)+1:].replace(os.path.sep,'.'))
+                elif nm[-3:] == '.py':
+                    stdlib_module_names.add(os.path.join(prefix, nm)[:-3].replace(os.path.sep,'.'))
+                elif nm[-3:] == '.so' and top[-11:] == 'lib-dynload':
+                    stdlib_module_names.add(nm[0:-3])
+        for builtin in sys.builtin_module_names:
+            stdlib_module_names.add(builtin)
+        return stdlib_module_names
+
+
+def _get_deps(path, include_stdlib=False):
+    stdlib_module_names = get_stdlib_module_names()
+    with open(path) as f:
+        root = ast.parse(f.read(), path)
+
+    deps = []
+    for node in ast.walk(root):
+        if isinstance(node, ast.Import):
+            for n in node.names:
+                mod = n.name.split('.')[0]
+                if include_stdlib or mod not in stdlib_module_names:
+                    deps.append(mod)
+        elif isinstance(node, ast.ImportFrom):
+            mod = node.module.split('.')[0]
+            if include_stdlib or mod not in stdlib_module_names:
+                deps.append(mod)
+        else:
+            continue
+
+    # allow comments to specify package name, e.g....
+    # import sklearn  # install scikit-learn
+    # TODO Make this robust a la https://stackoverflow.com/a/36055400/2903532
+    with open(path) as f:
+        install_lines = [line for line in f
+                         if 'import' in line
+                         and re.search(r'#\s+install\s+', line)]
+    dep_pkg_list = []
+    for dep in deps:
+        pkgs = [match.group(1).split()
+                for line in install_lines
+                if dep in line
+                and (match := re.search(r'#\s+install\s+(.+)$', line))]
+        pkgs = set(pkg for group in pkgs for pkg in group)  # flattened
+        dep_pkg_list.append((dep, tuple(pkgs) or None))
+    return dep_pkg_list
+
+
+def autoinstall(interactive=True, upgrade=False):
+    filename = inspect.stack()[1].filename
+    deps = _get_deps(filename)
+
+    if not interactive:
+        print(f'autoinstall is installing the following: {" ".join(deps)} ...')
+        # pip_inside.install(*deps, upgrade=upgrade)
+    else:
+        for mod, pkgs in deps:
+            # TODO: check if dependency is installed
+            print(f'{dep} is not installed.')
+            while True:
+                resp = input(f'Install {dep}? (Y)es / (N)o / (C)ustomize name  ')
+                if re.search(r'y(?:es)?', resp, flags=re.I):
+                    print(f'Installing {dep}...', file=sys.stderr)
+                    # pip_inside.install(dep, upgrade=upgrade)
+                    break
+                elif re.search(r'no?', resp, flags=re.I):
+                    print('Okay. Taking no action.')
+                    break
+                elif re.search(r'c(?:ustom)?', resp, flags=re.I):
+                    pkg_name = input(f'What package name would you like to install instead of {dep}? ')
+                    print(f'Installing {install_name}...', file=sys.stderr)
+                    # pip_inside.install(pkg_name, upgrade=upgrade)
+                    break
+                else:
+                    print('Invalid input. Type "y", "n", or "c" and press [enter].')
diff --git a/src/pip_inside/install.py b/src/pip_inside/install.py
new file mode 100644
index 0000000..3651f76
--- /dev/null
+++ b/src/pip_inside/install.py
@@ -0,0 +1,151 @@
+from glob import glob
+import os
+from pathlib import Path
+import re
+import shlex
+import subprocess
+import sys
+from warnings import warn
+
+try:
+    from pip._internal.commands.install import InstallCommand
+    from pip._vendor.distlib.database import DistributionPath
+    install_cmd = InstallCommand(name='dummy', summary='Provides parse_args.')
+except ModuleNotFoundError:
+    raise ModuleNotFoundError('Please install pip for the current '
+                              'interpreter: (%s).' % sys.executable)
+
+__all__ = ['install']
+
+
+pipfiles = []
+for i in range(4):  # TODO: 4 is completely arbitrary here. improve?
+    pipfiles.extend(glob(('..' + os.sep) * i + 'Pipfile'))
+if pipfiles:
+    pipfiles = [os.path.abspath(f) for f in pipfiles]
+    msg = ('Warning: the following Pipfiles will be bypassed by '
+           'pip_inside.install:\n\t' + '\n\t'.join(pipfiles))
+    warn(msg, stacklevel=2)
+
+
+def _parse_target(target):
+    """Parse install target down to the distribution name."""
+    # TODO: make more robust
+    if target.endswith('.whl'):
+        return Path(target).name.split('-')[0]
+    else:
+        # https://pip.pypa.io/en/stable/reference/requirement-specifiers/
+        return re.search(r'(.+?)(?:[!~<>=]{1,2}.+)?$', target).group(1)
+
+
+def install(*args, **kwargs):
+    """Install packages into the current environment.
+
+    Equivalent examples of command-line pip and pip_inside are grouped below.
+
+    METHOD 1: Single argument is exactly the same as command line interface,
+    beginning with 'pip install ...'
+
+    $ pip install --user --upgrade some_pkg
+    >>> install('pip install --user --upgrade some_pkg')
+
+    METHOD 2: Arguments from command-line implementation split on spaces
+
+    $ pip install some_pkg
+    >>> install('some_pkg')
+
+    $ pip install --user --upgrade some_pkg
+    >>> install('--user', '--upgrade', 'some_pkg')
+    >>> install(*'--user --upgrade some_pkg'.split())
+
+    If preferred, keyword-value arguments can also be used:
+
+    $ pip install -r requirements.txt
+    >>> install(r='requirements.txt')
+    >>> install('-r', 'requirements.txt')
+    >>> install(*'-r requirements.txt'.split())
+
+    $ pip install --no-index --find-links /local/dir/ some_pkg
+    # Note the use of '_' in the following keyword example.
+    >>> install('--no-index', 'some_pkg', find_links='/local/dir/')
+    >>> install('--no-index', '--find-links', '/local/dir/', 'some_pkg')
+    >>> install(*'--no-index --find-links /local/dir/ some_pkg'.split())
+
+    """
+    cli_args = _build_install_cmd(*args, **kwargs)
+    # use pip internals to isolate package names
+    _, req_targets = install_cmd.parse_args(cli_args)
+    assert req_targets[:2] == ['pip', 'install']
+    req_targets = [_parse_target(t) for t in req_targets[2:]]
+    dist_path = DistributionPath(include_egg=True)
+    target_providers = [d.modules
+                        for t in req_targets
+                        for d in dist_path.get_distributions()  # installed
+                        if t == d.name]
+    target_providers = [y for x in target_providers for y in x]  # flatten
+    target_origins = {}
+    for t in target_providers:
+        if t in sys.modules:
+            target_origins[sys.modules[t].__spec__.origin] = t
+    dist_path = DistributionPath(include_egg=True)
+    dists = [dist_path.get_distribution(t) for t in req_targets]
+    paths = set()
+    for dist in dists:
+        if dist is not None:
+            for path, _, _ in dist.list_installed_files():
+                paths.add(os.path.join(os.path.dirname(dist.path), path))
+    already_loaded = paths.intersection(target_origins)
+    already_loaded = [target_origins[origin] for origin in already_loaded]
+    print('Trying  ', ' '.join(cli_args), '  ...', file=sys.stderr)
+    cli_cmd = [sys.executable, "-m"] + cli_args
+    result = subprocess.run(cli_cmd, check=True)
+    if result.returncode == 0 and already_loaded:
+        warn('WARNING! The following modules were already loaded. Restart '
+             f'python to see changes:  {", ".join(already_loaded)}',
+             UserWarning)
+    return result
+
+
+def _build_install_cmd(*args, **kwargs):
+    if len(args) == 1 and args[0].startswith('pip install '):
+        return shlex.split(args[0])
+    else:
+        cli_args = ['pip', 'install']
+        # Keyword arguments are translated to CLI options
+        for raw_k, v in kwargs.items():
+            k = raw_k.replace('_', '-')  # Python identifiers -> CLI long names
+            append_value = isinstance(v, str)
+            if append_value:
+                # When arg value is str, append both it and option to CLI args
+                append_option = True
+                if not v:
+                    raise ValueError("Empty string passed as value for option "
+                                     "{}".format(k))
+            else:
+                # assume the value indicates whether to include a boolean flag.
+                # None->omit, true->include, false->include negated
+                append_option = v is not None
+                if k.startswith("no-"):
+                    # suggest `some-option=True` instead of
+                    # `no-some-option=False`
+                    raw_suffix = raw_k[3:]
+                    msg_template = ("Rather than '{}={!r}', "
+                                    "try '{{}}={{!r}}'".format(raw_k, v))
+                    if append_option:
+                        suggestion = msg_template.format(raw_suffix, not v)
+                    else:
+                        suggestion = msg_template.format(raw_suffix, None)
+                    raise ValueError(suggestion)
+                if append_option and not v:
+                    k = "no-" + k
+            if append_option:
+                if len(k) == 1:  # short flag
+                    option = "-" + k
+                else:  # long flag
+                    option = "--" + k
+                cli_args.append(option)
+                if append_value:
+                    cli_args.append(v)
+        # Positional arguments are passed directly as CLI arguments
+        cli_args += args
+    return cli_args


Test patch:
diff --git a/tests/data/imports.py b/tests/data/imports.py
new file mode 100644
index 0000000..d989abf
--- /dev/null
+++ b/tests/data/imports.py
@@ -0,0 +1,5 @@
+from bs4 import BeautifulSoup as bs  # install beautifulsoup4
+import libcst as cst
+from sklearn.ensemble import RandomForestClassifier  # install scikit-learn
+import requests as r
+import numpy, pandas
diff --git a/tests/functional/test_autoinstall.py b/tests/functional/test_autoinstall.py
new file mode 100644
index 0000000..86a6551
--- /dev/null
+++ b/tests/functional/test_autoinstall.py
@@ -0,0 +1,10 @@
+from pip_inside.autoinstall import _get_deps
+
+
+def test_get_deps():
+    assert _get_deps('tests/data/imports.py') == [('bs4', ('beautifulsoup4',)),
+                                                  ('libcst', None),
+                                                  ('sklearn', ('scikit-learn',)),
+                                                  ('requests', None),
+                                                  ('numpy', None),
+                                                  ('pandas', None)]
diff --git a/tests/functional/test_install.py b/tests/functional/test_install.py
index c6e5fd5..a7227ee 100644
--- a/tests/functional/test_install.py
+++ b/tests/functional/test_install.py
@@ -2,7 +2,7 @@
 
 import pytest
 
-from pip_inside import _build_install_cmd
+from pip_inside.install import _build_install_cmd
 from pip_inside import install
 
 TEST_PKG = 'realpython-reader'