Repository: timothycrosley/isort
Pull request #275 (2015-05-31 01:50:03)

Custom section support
Users can now define their own sections and ordering. For example

    import_heading_stdlib = Standard Library
    import_heading_thirdparty = Third Party
    import_heading_firstparty = First Party
    import_heading_django = Django
    import_heading_pandas = Pandas
    known_django = django
    known_pandas = pandas,numpy
    known_first_party = p24,p24.imports._VERSION
    sections = FUTURE,STDLIB,DJANGO,THIRDPARTY,PANDAS,FIRSTPARTY,LOCALFOLDER

would create two new sections with the specified known modules.

    # Standard Library
    import os
    import p24.imports._argparse as argparse
    import p24.imports._subprocess as subprocess
    import sys

    # Django
    from django.conf import settings
    from django.db import models

    # Third Party
    from bottle import Bottle, redirect, response, run

    # Pandas
    import numpy as np
    import pandas as pd

    # First Party
    import p24.imports._VERSION as VERSION
    import p24.shared.media_wiki_syntax as syntax


Patch:
diff --git a/.gitignore b/.gitignore
index 04c6f1b5..186e1599 100644
--- a/.gitignore
+++ b/.gitignore
@@ -16,6 +16,7 @@ develop-eggs
 lib
 lib64
 MANIFEST
+.eggs
 
 # Installer logs
 pip-log.txt
diff --git a/README.md b/README.md
index cf959abb..709ecea3 100644
--- a/README.md
+++ b/README.md
@@ -143,6 +143,7 @@ and puts them all at the top of the file grouped together by the type of import:
 - Current Python Project
 - Explicitly Local (. before import, as in: from . import x)
 - Custom Separate Sections (Defined by forced_separate list in configuration file)
+- Custom Sections (Defined by sections list in configuration file)
 
 Inside of each section the imports are sorted alphabetically. isort automatically removes duplicate python imports,
 and wraps long from imports to the specified line length (defaults to 80).
@@ -286,6 +287,26 @@ Will be produced instead of:
 
 To enable this set 'balanced_wrapping' to True in your config or pass the -e option into the command line utility.
 
+Custom Sections and Ordering
+============================
+
+You can change the section order with `sections` option from the default of:
+
+    FUTURE,STDLIB,THIRDPARTY,FIRSTPARTY,LOCALFOLDER
+
+to your preference:
+
+    sections=FUTURE,STDLIB,FIRSTPARTY,THIRDPARTY,LOCALFOLDER
+
+You also can define your own sections and thier order.
+
+Example:
+
+    known_django=django
+    known_pandas=pandas,numpy
+    sections=FUTURE,STDLIB,DJANGO,THIRDPARTY,PANDAS,FIRSTPARTY,LOCALFOLDER
+
+would create two new sections with the specified known modules.
 
 Auto-comment import sections
 ======================
diff --git a/isort/__init__.py b/isort/__init__.py
index 4a445fa7..2c399d8e 100644
--- a/isort/__init__.py
+++ b/isort/__init__.py
@@ -23,6 +23,6 @@ OTHER DEALINGS IN THE SOFTWARE.
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 from . import settings
-from .isort import SECTION_NAMES, SECTIONS, SortImports
+from .isort import SortImports
 
 __version__ = "3.9.6"
diff --git a/isort/isort.py b/isort/isort.py
index b838b3e4..c2b3f8d0 100644
--- a/isort/isort.py
+++ b/isort/isort.py
@@ -41,8 +41,12 @@ from pies.overrides import *
 
 from . import settings
 
-SECTION_NAMES = ("FUTURE", "STDLIB", "THIRDPARTY", "FIRSTPARTY", "LOCALFOLDER")
-SECTIONS = namedtuple('Sections', SECTION_NAMES)(*range(len(SECTION_NAMES)))
+KNOWN_SECTION_MAPPING = {
+    'STDLIB': 'STANDARD_LIBRARY',
+    'FUTURE': 'FUTURE_LIBRARY',
+    'FIRSTPARTY': 'FIRST_PARTY',
+    'THIRDPARTY': 'THIRD_PARTY',
+}
 
 
 class SortImports(object):
@@ -59,7 +63,8 @@ class SortImports(object):
         self.config = settings.from_path(settings_path).copy()
         for key, value in itemsview(setting_overrides):
             access_key = key.replace('not_', '').lower()
-            if type(self.config.get(access_key)) in (list, tuple):
+            # The sections config needs to retain order and can't be converted to a set.
+            if access_key != 'sections' and type(self.config.get(access_key)) in (list, tuple):
                 if key.startswith('not_'):
                     self.config[access_key] = list(set(self.config[access_key]).difference(value))
                 else:
@@ -112,7 +117,10 @@ class SortImports(object):
         self.comments = {'from': {}, 'straight': {}, 'nested': {}, 'above': {'straight': {}, 'from': {}}}
         self.imports = {}
         self.as_map = {}
-        for section in itertools.chain(SECTIONS, self.config['forced_separate']):
+
+        section_names = self.config.get('sections')
+        self.sections = namedtuple('Sections', section_names)(*[n for n in section_names])
+        for section in itertools.chain(self.sections, self.config['forced_separate']):
             self.imports[section] = {'straight': set(), 'from': {}}
 
         self.index = 0
@@ -203,19 +211,16 @@ class SortImports(object):
                 return forced_separate
 
         if moduleName.startswith("."):
-            return SECTIONS.LOCALFOLDER
+            return self.sections.LOCALFOLDER
 
         # Try to find most specific placement instruction match (if any)
         parts = moduleName.split('.')
         module_names_to_check = ['.'.join(parts[:first_k]) for first_k in range(len(parts), 0, -1)]
         for module_name_to_check in module_names_to_check:
-            for placement, config_key in (
-                    (SECTIONS.FUTURE, 'known_future_library'),
-                    (SECTIONS.STDLIB, 'known_standard_library'),
-                    (SECTIONS.THIRDPARTY, 'known_third_party'),
-                    (SECTIONS.FIRSTPARTY, 'known_first_party'),
-                    ):
-                if module_name_to_check in self.config[config_key]:
+            for placement in reversed(self.sections):
+                known_placement = KNOWN_SECTION_MAPPING.get(placement, placement)
+                config_key = 'known_{0}'.format(known_placement.lower())
+                if module_name_to_check in self.config.get(config_key, []):
                     return placement
 
         paths = PYTHONPATH
@@ -231,13 +236,13 @@ class SortImports(object):
             if (os.path.exists(module_path + ".py") or os.path.exists(module_path + ".so") or
                (os.path.exists(package_path) and os.path.isdir(package_path))):
                 if "site-packages" in prefix or "dist-packages" in prefix:
-                    return SECTIONS.THIRDPARTY
+                    return self.sections.THIRDPARTY
                 elif "python2" in prefix.lower() or "python3" in prefix.lower():
-                    return SECTIONS.STDLIB
+                    return self.sections.STDLIB
                 else:
-                    return SECTIONS.FIRSTPARTY
+                    return self.sections.FIRSTPARTY
 
-        return SECTION_NAMES.index(self.config['default_section'])
+        return self.config['default_section']
 
     def _get_line(self):
         """Returns the current line from the file while incrementing the index."""
@@ -425,7 +430,7 @@ class SortImports(object):
 
         """
         output = []
-        for section in itertools.chain(SECTIONS, self.config['forced_separate']):
+        for section in itertools.chain(self.sections, self.config['forced_separate']):
             straight_modules = list(self.imports[section]['straight'])
             straight_modules = natsorted(straight_modules, key=lambda key: self._module_key(key, self.config))
             from_modules = sorted(list(self.imports[section]['from'].keys()))
@@ -441,8 +446,6 @@ class SortImports(object):
 
             if section_output:
                 section_name = section
-                if section in SECTIONS:
-                    section_name = SECTION_NAMES[section]
                 if section_name in self.place_imports:
                     self.place_imports[section_name] = section_output
                     continue
diff --git a/isort/main.py b/isort/main.py
index c2b42cf2..ff26882e 100755
--- a/isort/main.py
+++ b/isort/main.py
@@ -28,8 +28,8 @@ import sys
 import setuptools
 from pies.overrides import *
 
-from isort import SECTION_NAMES, SortImports, __version__
-from isort.settings import default, from_path
+from isort import SortImports, __version__
+from isort.settings import DEFAULT_SECTIONS, default, from_path
 
 
 def iter_source_code(paths):
@@ -143,7 +143,7 @@ def create_parser():
                         help='Forces all from imports to appear on their own line')
     parser.add_argument('-sd', '--section-default', dest='default_section',
                         help='Sets the default section for imports (by default FIRSTPARTY) options: ' +
-                        str(SECTION_NAMES))
+                        str(DEFAULT_SECTIONS))
     parser.add_argument('-df', '--diff', dest='show_diff', default=False, action='store_true',
                         help="Prints a diff of all the changes isort would make to a file, instead of "
                              "changing it in place")
diff --git a/isort/settings.py b/isort/settings.py
index 8e44611a..57067f61 100644
--- a/isort/settings.py
+++ b/isort/settings.py
@@ -36,6 +36,7 @@ except ImportError:
     import ConfigParser as configparser
 
 MAX_CONFIG_SEARCH_DEPTH = 25 # The number of parent directories isort will look for a config file within
+DEFAULT_SECTIONS = ("FUTURE", "STDLIB", "THIRDPARTY", "FIRSTPARTY", "LOCALFOLDER")
 
 WrapModes = ('GRID', 'VERTICAL', 'HANGING_INDENT', 'VERTICAL_HANGING_INDENT', 'VERTICAL_GRID', 'VERTICAL_GRID_GROUPED')
 WrapModes = namedtuple('WrapModes', WrapModes)(*range(len(WrapModes)))
@@ -45,6 +46,7 @@ default = {'force_to_top': [],
            'skip': ['__init__.py', ],
            'line_length': 79,
            'wrap_length': 0,
+           'sections': DEFAULT_SECTIONS,
            'known_future_library': ['__future__'],
            'known_standard_library': ["abc", "anydbm", "argparse", "array", "asynchat", "asyncore", "atexit", "base64",
                                       "BaseHTTPServer", "bisect", "bz2", "calendar", "cgitb", "cmd", "codecs",
diff --git a/setup.cfg b/setup.cfg
index 5e409001..4ab7eb2b 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,2 +1,6 @@
 [wheel]
 universal = 1
+
+[flake8]
+ignore = F401,F403,E502,E123,E127,E128,E303,E713,E111,E241,E302,E121,E261,W391
+max-line-length = 160


Test patch:
diff --git a/test_isort.py b/test_isort.py
index d8467498..41458991 100644
--- a/test_isort.py
+++ b/test_isort.py
@@ -665,6 +665,27 @@ def test_default_section():
                                   "\n"
                                   "import django.settings\n")
 
+def test_first_party_overrides_standard_section():
+    """Test to ensure changing the default section works as expected."""
+    test_input = ("import sys\n"
+                  "import os\n"
+                  "import profile.test\n")
+    test_output = SortImports(file_contents=test_input, known_first_party=['profile']).output
+    assert test_output == ("import os\n"
+                           "import sys\n"
+                           "\n"
+                           "import profile.test\n")
+
+def test_thirdy_party_overrides_standard_section():
+    """Test to ensure changing the default section works as expected."""
+    test_input = ("import sys\n"
+                  "import os\n"
+                  "import profile.test\n")
+    test_output = SortImports(file_contents=test_input, known_third_party=['profile']).output
+    assert test_output == ("import os\n"
+                           "import sys\n"
+                           "\n"
+                           "import profile.test\n")
 
 def test_force_single_line_imports():
     """Test to ensure forcing imports to each have their own line works as expected."""
@@ -1172,7 +1193,6 @@ def test_place_comments():
                            "import os\n"
                            "import sys\n")
 
-
 def test_placement_control():
     """Ensure that most specific placement control match wins"""
     test_input = ("import os\n"
@@ -1187,6 +1207,7 @@ def test_placement_control():
                 known_standard_library=['p24.imports'],
                 known_third_party=['bottle'],
                 default_section="THIRDPARTY").output
+
     assert test_output == ("import os\n"
                            "import p24.imports._argparse as argparse\n"
                            "import p24.imports._subprocess as subprocess\n"
@@ -1198,6 +1219,54 @@ def test_placement_control():
                            "import p24.shared.media_wiki_syntax as syntax\n")
 
 
+def test_custom_sections():
+    """Ensure that most specific placement control match wins"""
+    test_input = ("import os\n"
+                  "import sys\n"
+                  "from django.conf import settings\n"
+                  "from bottle import Bottle, redirect, response, run\n"
+                  "import p24.imports._argparse as argparse\n"
+                  "from django.db import models\n"
+                  "import p24.imports._subprocess as subprocess\n"
+                  "import pandas as pd\n"
+                  "import p24.imports._VERSION as VERSION\n"
+                  "import numpy as np\n"
+                  "import p24.shared.media_wiki_syntax as syntax\n")
+    test_output = SortImports(file_contents=test_input,
+                known_first_party=['p24', 'p24.imports._VERSION'],
+                import_heading_stdlib='Standard Library',
+                import_heading_thirdparty='Third Party',
+                import_heading_firstparty='First Party',
+                import_heading_django='Django',
+                import_heading_pandas='Pandas',
+                known_standard_library=['p24.imports'],
+                known_third_party=['bottle'],
+                known_django=['django'],
+                known_pandas=['pandas', 'numpy'],
+                default_section="THIRDPARTY",
+                sections=["FUTURE", "STDLIB", "DJANGO", "THIRDPARTY", "PANDAS", "FIRSTPARTY", "LOCALFOLDER"]).output
+    assert test_output == ("# Standard Library\n"
+                           "import os\n"
+                           "import p24.imports._argparse as argparse\n"
+                           "import p24.imports._subprocess as subprocess\n"
+                           "import sys\n"
+                           "\n"
+                           "# Django\n"
+                           "from django.conf import settings\n"
+                           "from django.db import models\n"
+                           "\n"
+                           "# Third Party\n"
+                           "from bottle import Bottle, redirect, response, run\n"
+                           "\n"
+                           "# Pandas\n"
+                           "import numpy as np\n"
+                           "import pandas as pd\n"
+                           "\n"
+                           "# First Party\n"
+                           "import p24.imports._VERSION as VERSION\n"
+                           "import p24.shared.media_wiki_syntax as syntax\n")
+
+
 def test_sticky_comments():
     """Test to ensure it is possible to make comments 'stick' above imports"""
     test_input = ("import os\n"
@@ -1328,4 +1397,3 @@ def test_fcntl():
                   "import os\n"
                   "import sys\n")
     assert SortImports(file_contents=test_input).output == test_input
-