Merge pull request #286 from timothycrosley/custom-sections

Custom sections

8 files changed, 126 insertions, 25 deletions

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 2a13aa30..8f8ddaa0 100644
--- a/isort/isort.py
+++ b/isort/isort.py
@@ -42,8 +42,12 @@ from .pie_slice 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):
@@ -60,7 +64,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:
@@ -113,7 +118,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
@@ -204,19 +212,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
@@ -232,13 +237,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."""
@@ -430,7 +435,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 = nsorted(straight_modules, key=lambda key: self._module_key(key, self.config))
             from_modules = sorted(list(self.imports[section]['from'].keys()))
@@ -446,8 +451,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 6e3188aa..86270675 100755
--- a/isort/main.py
+++ b/isort/main.py
@@ -28,8 +28,8 @@ import sys
 import setuptools
 from .pie_slice 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 e71833cc..a14fa115 100644
--- a/isort/settings.py
+++ b/isort/settings.py
@@ -35,6 +35,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)))
@@ -44,6 +45,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
diff --git a/test_isort.py b/test_isort.py
index a9f2adc8..574f201e 100644
--- a/test_isort.py
+++ b/test_isort.py
@@ -677,6 +677,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."""
@@ -1184,7 +1205,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"
@@ -1199,6 +1219,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"
@@ -1210,6 +1231,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"
@@ -1341,8 +1410,9 @@ def test_fcntl():
                   "import sys\n")
     assert SortImports(file_contents=test_input).output == test_input
 
-def test_import_split_is_word_boundary_aware():
 
+def test_import_split_is_word_boundary_aware():
+    """Test to ensure that isort splits words in a boundry aware mannor"""
     test_input = ("from mycompany.model.size_value_array_import_func import ("
                 "    get_size_value_array_import_func_jobs,"
                 ")")