Merge pull request #329 from Carreau/resort-imports

Custom implementation of import sorting.

bin/tidy-imports

@@ -21,13 +21,17 @@ Only top-level import statements are touched.
 # License: MIT http://opensource.org/licenses/MIT
 
 
+from __future__ import print_function
+
 import os
 
 from   pyflyby._cmdline         import hfmt, parse_args, process_actions
-from pyflyby._imports2s import (canonicalize_imports,
-                                fix_unused_and_missing_imports,
-                                replace_star_imports,
-                                transform_imports)
+from   pyflyby._import_sorting  import sort_imports
+from   pyflyby._imports2s       import (canonicalize_imports,
+                                        fix_unused_and_missing_imports,
+                                        replace_star_imports,
+                                        transform_imports)
+from   pyflyby._parse           import PythonBlock
 
 import toml
 TOML_AVAIL = True
@@ -104,6 +108,10 @@ def _addopts(parser):
                         help=hfmt('''
                             (Default) Replace imports with canonical
                             equivalent imports, according to database.'''))
+    parser.add_option('--experimental-sort-imports',
+                        default=False, action='store_true',
+                        help=hfmt('''
+                            experimental import sorting'''))
     parser.add_option('--no-canonicalize', dest='canonicalize',
                         default=True, action='store_false',
                         help=hfmt('''
@@ -128,7 +136,7 @@ def _addopts(parser):
                         help=hfmt('''
                             Equivalent to --no-add-missing
                             --no-add-mandatory.'''))
-def main():
+def main() -> None:
 
     config_text = _get_pyproj_toml_config()
     if config_text:
@@ -139,28 +147,43 @@ def main():
     def _add_opts_and_defaults(parser):
         _addopts(parser)
         parser.set_defaults(**default_config)
+
     options, args = parse_args(
         _add_opts_and_defaults, import_format_params=True, modify_action_params=True, defaults=default_config)
 
-    def modify(x):
+    def modify(block:PythonBlock) -> PythonBlock:
         if options.transformations:
-            x = transform_imports(x, options.transformations,
+            block = transform_imports(block, options.transformations,
                                   params=options.params)
         if options.replace_star_imports:
-            x = replace_star_imports(x, params=options.params)
-        x = fix_unused_and_missing_imports(
-            x, params=options.params,
+            block = replace_star_imports(block, params=options.params)
+        block = fix_unused_and_missing_imports(
+            block, params=options.params,
             add_missing=options.add_missing,
             remove_unused=options.remove_unused,
             add_mandatory=options.add_mandatory,
             )
         # TODO: disable sorting until we figure out #287
         # https://github.com/deshaw/pyflyby/issues/287
-        # from pyflyby._imports2s import sort_imports
-        # x = sort_imports(x)
+
+        # here we get a (single?) PythonBlock, we can access each statement with 
+        # >>> block.statements
+        # and each statement can have a:
+        # is_import
+        # or
+        # is_comment or blank
+
+        # TODO: we do Python(str(...)) in order to unparse-reparse and get proper ast node numbers.
+        if options.experimental_sort_imports:
+            sorted_imports = PythonBlock(str(sort_imports(block)))
+        else:
+            sorted_imports = block
         if options.canonicalize:
-            x = canonicalize_imports(x, params=options.params)
-        return x
+            cannonical_imports = canonicalize_imports(sorted_imports, params=options.params)
+        else:
+            cannonical_imports = sorted_imports
+        return cannonical_imports
+
     process_actions(args, options.actions, modify)
 
 

lib/python/pyflyby/_import_sorting.py

@@ -0,0 +1,165 @@
+"""
+This module contain utility functions to sort imports in a Python Block.
+"""
+from __future__ import annotations
+
+from collections import Counter
+from dataclasses import dataclass
+from itertools import groupby
+from pyflyby._importstmt import ImportStatement, PythonStatement
+from pyflyby._parse import PythonBlock
+from typing import List, Tuple, Union
+
+
+# @dataclass
+# class ImportSection:
+#     """
+#     This represent an Import
+#     """
+#     sections: List[ImportGroup]
+
+
+@dataclass
+class ImportGroup:
+    """
+    Typically at the top of a file, the first import will be part
+    of an ImportGroup, and subsequent imports
+    will be other import groups which.
+
+    Import sorting will affect only the imports in the same group,
+    as we want import sorting to be indempotent. If import sorting
+    were migrating imports between groups, then imports could be
+    moved to sections were comments would nto be relevant.
+    """
+
+    imports: List[ImportStatement]
+
+    @classmethod
+    def from_statements(cls, statements: List[PythonStatement]) -> ImportGroup:
+        return ImportGroup([ImportStatement(s) for s in statements])
+
+    def sorted(self) -> ImportGroup:
+        """
+        return an ImportGroup with import sorted lexicographically.
+        """
+        return ImportGroup(sorted(self.imports, key=lambda x: x._cmp()))
+
+    def sorted_subgroups(self) -> List[Tuple[bool, List[ImportStatement]]]:
+        """
+        Return a list of subgroup keyed by module and whether they are the sole import
+        from this module.
+
+        From issue #13, we will want to sort import from the same package together when
+        there is more than one import and separat with a blank line
+
+        We also group all imports from a package that appear only once together.
+
+        For this we need both to know if the import is from a single import (the boolean).
+
+        Returns
+        -------
+        bool:
+           wether from a single import
+        List[ImportStatement]
+           The actual import.
+        """
+        c = Counter(imp.module[0] for imp in self.imports)
+        return [
+            (c[k] > 1, list(v)) for k, v in groupby(self.imports, lambda x: x.module[0])
+        ]
+
+
+def split_import_groups(
+    statements: Tuple[PythonStatement],
+) -> List[Union[ImportGroup, PythonStatement]]:
+    """
+    Given a list of statements split into import groups.
+
+    One of the question is how to treat split with comments.
+
+     - Do blank lines after groups with comments start a new block
+     - Does the comment line create a complete new block.
+
+    In particular because we want this to be indempotent,
+    we can't move imports between groups.
+    """
+
+    # these are the import groups we'll
+    groups: List[PythonStatement | ImportGroup] = []
+
+    current_group: List[PythonStatement] = []
+    statemt_iterator = iter(statements)
+    for statement in statemt_iterator:
+        if statement.is_blank and current_group:
+            pass
+            # currently do nothing with whitespace while in import groups.
+        elif statement.is_import:
+            # push on top of current_comment.
+            current_group.append(statement)
+        else:
+            if current_group:
+                groups.append(ImportGroup.from_statements(current_group))
+            current_group = []
+            groups.append(statement)
+
+            # we should break of and populate rest
+            # We can't do anything if we encounter any non-import statement,
+            # as we do no know if it can be a conditional.
+            # technically I guess we coudl find another import block, and reorder these.
+    if current_group:
+        groups.append(ImportGroup.from_statements(current_group))
+
+    # this is an iterator, not an iterable, we exaust it to reify the rest.
+
+    # first group may be empty if the first line is a comment.
+    # We filter and sort relevant statements.
+    groups = [g for g in groups if groups]
+    sorted_groups: List[PythonStatement | ImportGroup] = [
+        g.sorted() if isinstance(g, ImportGroup) else g for g in groups
+    ]
+    return sorted_groups
+
+
+def regroup(groups: List[ImportGroup | PythonStatement]) -> PythonBlock:
+    """
+    given import groups and list of statement, return an Python block with sorted import
+    """
+    res: str = ""
+    in_single = False
+    for group in groups:
+        if isinstance(group, ImportGroup):
+            # the subgroup here will be responsible for  groups reordering.
+            for mult, subgroup in group.sorted_subgroups():
+                if mult:
+                    if in_single:
+                        res += "\n"
+                    if not res.endswith("\n"):
+                        res += "\n"
+                    for x in subgroup:
+                        res += x.pretty_print(import_column=30).rstrip() + "\n"
+                    if not res.endswith("\n\n"):
+                        res += "\n"
+                    in_single = False
+                else:
+                    assert len(subgroup) == 1
+                    in_single = True
+                    for sub in subgroup:
+                        res += str(sub).rstrip() + "\n"
+
+            if not res.endswith("\n\n"):
+                res += "\n"
+        else:
+            if in_single and not res.endswith("\n\n"):
+                res += "\n"
+            in_single = False
+            res += str(PythonBlock(group))
+
+    return PythonBlock.concatenate([PythonBlock(res)])
+
+
+def sort_imports(block: PythonBlock) -> PythonBlock:
+    assert isinstance(block, PythonBlock)
+    # we ignore below that block.statement can be a List[Unknown]
+    gs = split_import_groups(block.statements)  # type: ignore
+    # TODO: math the other fileds like filename....
+    return regroup(gs)