Match on module aliases for auto import suggestions

rope/base/prefs.py

@@ -140,6 +140,22 @@ class Prefs:
         """),
     )
 
+    import_aliases: List[Tuple[str, str]] = field(
+        default_factory=lambda : [
+            ("np", "numpy"),
+            ("pd", "pandas"),
+            ("plt", "matplotlib.pyplot"),
+            ("sns", "seaborn"),
+            ("tf", "tensorflow"),
+            ("sk", "sklearn"),
+            ("sm", "statsmodels"),
+        ],
+        description=dedent("""
+            Aliases for module names.  For example, `[('np', 'numpy')]` makes rope recommend 
+            ``import numpy as np``.
+        """),
+    )
+
     prefer_module_from_imports: bool = field(
         default=False,
         description=dedent("""

rope/contrib/autoimport/defs.py

@@ -92,6 +92,12 @@ class Package(NamedTuple):
     type: PackageType
 
 
+class Alias(NamedTuple):
+    """A module alias to be added to the database."""
+    alias: str
+    modname: str
+
+
 class Name(NamedTuple):
     """A Name to be added to the database."""
 

rope/contrib/autoimport/models.py

@@ -85,6 +85,27 @@ class Metadata(Model):
     objects = Query(table_name, columns)
 
 
+class Alias(Model):
+    table_name = "aliases"
+    schema = {
+        "alias": "TEXT",
+        "module": "TEXT",
+    }
+    columns = list(schema.keys())
+    objects = Query(table_name, columns)
+
+    @classmethod
+    def create_table(cls, connection):
+        super().create_table(connection)
+        connection.execute("CREATE INDEX IF NOT EXISTS alias ON aliases(alias)")
+
+    modules = Query(
+        "(SELECT DISTINCT aliases.*, package, source, type FROM aliases INNER JOIN names on aliases.module = names.module)",
+        columns + ["package", "source", "type"],
+    )
+    search_modules_with_alias = modules.where("alias LIKE (?)")
+
+
 class Name(Model):
     table_name = "names"
     schema = {

rope/contrib/autoimport/sqlite.py

@@ -19,6 +19,7 @@
 from rope.contrib.autoimport import models
 from rope.contrib.autoimport.defs import (
     ModuleFile,
+    Alias,
     Name,
     NameType,
     Package,
@@ -293,6 +294,12 @@ def _search_module(
             yield SearchResult(
                 f"import {module}", module, source, NameType.Module.value
             )
+        for alias, module, source in self._execute(
+            models.Alias.search_modules_with_alias.select("alias", "module", "source"), (name,)
+        ):
+            yield SearchResult(
+                f"import {module} as {alias}", alias, source, NameType.Module.value
+            )
 
     def get_modules(self, name) -> List[str]:
         """Get the list of modules that have global `name`."""
@@ -434,11 +441,14 @@ def clear_cache(self):
         """
         with self.connection:
             self._execute(models.Name.objects.drop_table())
+            self._execute(models.Alias.objects.drop_table())
             self._execute(models.Package.objects.drop_table())
             self._execute(models.Metadata.objects.drop_table())
             models.Name.create_table(self.connection)
+            models.Alias.create_table(self.connection)
             models.Package.create_table(self.connection)
             models.Metadata.create_table(self.connection)
+            self.add_aliases(self.project.prefs.import_aliases)
             data = (
                 versioning.calculate_version_hash(self.project),
                 json.dumps(versioning.get_version_hash_data(self.project)),
@@ -548,6 +558,10 @@ def _convert_name(name: Name) -> tuple:
             name.source.value,
             name.name_type.value,
         )
+
+    def add_aliases(self, aliases: Iterable[Alias]):
+        if aliases:
+            self._executemany(models.Alias.objects.insert_into(), aliases)
 
     def _add_names(self, names: Iterable[Name]):
         if names is not None:

ropetest/contrib/autoimporttest.py

@@ -1,6 +1,7 @@
 import unittest
 
 from rope.contrib.autoimport import sqlite as autoimport
+from rope.contrib.autoimport.defs import Alias
 from ropetest import testutils
 
 
@@ -124,6 +125,30 @@ def test_search_module(self):
         self.assertIn(import_statement, self.importer.search("os"))
         self.assertIn(import_statement, self.importer.search("o"))
 
+    def test_search_alias(self):
+        self.mod2.write("myvar = None\n")
+        self.importer.update_resource(self.mod2)
+        self.importer.add_aliases([
+            ("noMatch", "does_not_exists_this"),
+            ("hasMatch", "pkg.mod2"),
+        ])
+
+        self.assertEqual([], self.importer.search("noMatch", exact_match=True))
+
+        import_statement = ("import pkg.mod2 as hasMatch", "hasMatch")
+        self.assertIn(import_statement, self.importer.search("hasMatch", exact_match=True))
+        self.assertIn(import_statement, self.importer.search("hasM"))
+        self.assertIn(import_statement, self.importer.search("h"))
+
+    def test_alias_updated_from_prefs(self):
+        self.mod2.write("myvar = None\n")
+        self.project.prefs.import_aliases = [("mod2_alias", "pkg.mod2")]
+        self.importer.clear_cache()
+        self.importer.update_resource(self.mod2)
+        import_statement = ("import pkg.mod2 as mod2_alias", "mod2_alias")
+        self.assertIn(import_statement, self.importer.search("mod2_alias", exact_match=True))
+        self.assertIn(import_statement, self.importer.search("mod2", exact_match=False))
+
     def test_search(self):
         self.importer.update_module("typing")
         import_statement = ("from typing import Dict", "Dict")

ropetest/testutils.py

@@ -26,6 +26,7 @@ def sample_project(foldername=None, **kwds):
         "automatic_soa": False,
         "ignored_resources": [".ropeproject", "*.pyc"],
         "import_dynload_stdmods": False,
+        "import_aliases": [],
     }
     prefs.update(kwds)
     project = rope.base.project.Project(root, **prefs)