Repository: miki725/importanize

Issue:
comments in import produce invalid syntax
```
from a import b, c  # noqa
```

is transformed to:

```
from a import (
    b,
    c  # noqa,
)
```

which is of course invalid syntax

Discussion:

Patch:
diff --git a/HISTORY.rst b/HISTORY.rst
index 2a6e41e..452409e 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -3,6 +3,13 @@
 History
 -------
 
+0.3 (2015-01-18)
+~~~~~~~~~~~~~~~~
+
+* Using tokens to parse Python files. As a result this allows to
+  fix how comments are handled
+  (see `#21 <https://github.com/miki725/importanize/issues/21>`_ for example)
+
 0.2 (2014-10-30)
 ~~~~~~~~~~~~~~~~
 
diff --git a/Makefile b/Makefile
index ef86e42..db3d1bc 100644
--- a/Makefile
+++ b/Makefile
@@ -1,7 +1,7 @@
 .PHONY: clean-pyc clean-build docs clean
 
 NOSE_FLAGS=-sv --with-doctest --rednose
-COVER_CONFIG_FLAGS=--with-coverage --cover-package=importanize,tests --cover-tests
+COVER_CONFIG_FLAGS=--with-coverage --cover-package=importanize,tests --cover-tests --cover-erase
 COVER_REPORT_FLAGS=--cover-html --cover-html-dir=htmlcov
 COVER_FLAGS=${COVER_CONFIG_FLAGS} ${COVER_REPORT_FLAGS}
 
diff --git a/importanize/__init__.py b/importanize/__init__.py
index 5e712ad..f1714eb 100755
--- a/importanize/__init__.py
+++ b/importanize/__init__.py
@@ -4,4 +4,4 @@ from __future__ import print_function, unicode_literals
 
 __author__ = 'Miroslav Shubernetskiy'
 __email__ = 'miroslav@miki725.com'
-__version__ = '0.2'
+__version__ = '0.3'
diff --git a/importanize/groups.py b/importanize/groups.py
index 98e9d98..7224ce1 100644
--- a/importanize/groups.py
+++ b/importanize/groups.py
@@ -16,7 +16,7 @@ class BaseImportGroup(object):
         self.config = config or {}
 
         self.statements = []
-        self.artifacts = kwargs.get('artifacts', {})
+        self.file_artifacts = kwargs.get('file_artifacts', {})
 
     @property
     def unique_statements(self):
@@ -58,12 +58,12 @@ class BaseImportGroup(object):
         return False
 
     def as_string(self):
-        sep = self.artifacts.get('sep', '\n')
+        sep = self.file_artifacts.get('sep', '\n')
         return sep.join(map(operator.methodcaller('as_string'),
                             self.unique_statements))
 
     def formatted(self):
-        sep = self.artifacts.get('sep', '\n')
+        sep = self.file_artifacts.get('sep', '\n')
         return sep.join(map(operator.methodcaller('formatted'),
                             self.unique_statements))
 
@@ -111,7 +111,7 @@ GROUP_MAPPING = OrderedDict((
 class ImportGroups(object):
     def __init__(self, **kwargs):
         self.groups = []
-        self.artifacts = kwargs.get('artifacts', {})
+        self.file_artifacts = kwargs.get('file_artifacts', {})
 
     def all_line_numbers(self):
         return sorted(list(set(list(
@@ -151,14 +151,14 @@ class ImportGroups(object):
             raise ValueError(msg)
 
     def as_string(self):
-        sep = self.artifacts.get('sep', '\n') * 2
+        sep = self.file_artifacts.get('sep', '\n') * 2
         return sep.join(filter(
             None, map(operator.methodcaller('as_string'),
                       self.groups)
         ))
 
     def formatted(self):
-        sep = self.artifacts.get('sep', '\n') * 2
+        sep = self.file_artifacts.get('sep', '\n') * 2
         return sep.join(filter(
             None, map(operator.methodcaller('formatted'),
                       self.groups)
diff --git a/importanize/main.py b/importanize/main.py
index bd02d60..0c239fa 100644
--- a/importanize/main.py
+++ b/importanize/main.py
@@ -12,7 +12,11 @@ import six
 
 from . import __version__
 from .groups import ImportGroups
-from .parser import find_imports_from_lines, get_artifacts, parse_statements
+from .parser import (
+    find_imports_from_lines,
+    get_file_artifacts,
+    parse_statements,
+)
 from .utils import read
 
 
@@ -73,16 +77,19 @@ parser.add_argument(
     nargs='*',
     default=['.'],
     help='Path either to a file or directory where '
-         'all Python import will be organized. ',
+         'all Python files imports will be organized.',
 )
 parser.add_argument(
     '-c', '--config',
     type=six.text_type,
     default=default_config,
     help='Path to importanize config json file. '
-         'If importanize.json is present, that config '
+         'If "{}" is present in either current folder '
+         'or any parent folder, that config '
          'will be used. Otherwise crude default pep8 '
-         'config will be used.{}'.format(found_default),
+         'config will be used.{}'
+         ''.format(IMPORTANIZE_CONFIG,
+                   found_default),
 )
 parser.add_argument(
     '--print',
@@ -113,7 +120,7 @@ def run_importanize(path, config, args):
             return False
 
     text = read(path)
-    artifacts = get_artifacts(path)
+    file_artifacts = get_file_artifacts(path)
 
     lines_iterator = enumerate(iter(text.splitlines()))
     imports = list(parse_statements(find_imports_from_lines(lines_iterator)))
@@ -149,10 +156,10 @@ def run_importanize(path, config, args):
              + lines[first_import_line_number:]
              + [''])
 
-    lines = artifacts.get('sep', '\n').join(lines)
+    lines = file_artifacts.get('sep', '\n').join(lines)
 
     if args.print:
-        print(lines.encode('utf-8'))
+        print(lines.encode('utf-8') if not six.PY3 else lines)
     else:
         with open(path, 'wb') as fid:
             fid.write(lines.encode('utf-8'))
diff --git a/importanize/parser.py b/importanize/parser.py
index 87df4c9..8763477 100644
--- a/importanize/parser.py
+++ b/importanize/parser.py
@@ -2,14 +2,41 @@
 from __future__ import print_function, unicode_literals
 import re
 
+import six
+
 from .statements import DOTS, ImportLeaf, ImportStatement
-from .utils import list_strip, read
+from .utils import list_split, read
+
+
+STATEMENT_COMMENTS = ('noqa',)
+TOKEN_REGEX = re.compile(r' +|[\(\)]|([,\\])|(#.*)')
+SPECIAL_TOKENS = (',', 'import', 'from', 'as')
+
 
+class Token(six.text_type):
+    def __new__(cls, value, **kwargs):
+        obj = six.text_type.__new__(cls, value)
+        obj.is_comment_first = False
+        for k, v in kwargs.items():
+            setattr(obj, k, v)
+        return obj
 
-CHARS = re.compile(r'[\\()]')
+    @property
+    def is_comment(self):
+        return self.startswith('#')
+
+    @property
+    def normalized(self):
+        if not self.is_comment:
+            return self
+        else:
+            if self.startswith('# '):
+                return self[2:]
+            else:
+                return self[1:]
 
 
-def get_artifacts(path):
+def get_file_artifacts(path):
     """
     Get artifacts for the given file.
 
@@ -37,10 +64,7 @@ def get_artifacts(path):
 
 def find_imports_from_lines(iterator):
     """
-    Find import statements as strings from
-    enumerated iterator of lines.
-
-    Usually the iterator will of file lines.
+    Find only import statements from enumerated iterator of file files.
 
     Parameters
     ----------
@@ -51,28 +75,9 @@ def find_imports_from_lines(iterator):
     Returns
     -------
     imports : generator
-        Iterator which yields normalized import
-        strings with line numbers from which the
-        import statement has been parsed from.
-
-    Example
-    -------
-
-    ::
-
-        >>> parsed = list(find_imports_from_lines(iter([
-        ...     (1, 'from __future__ import unicode_literals'),
-        ...     (2, 'import os.path'),
-        ...     (3, 'from datetime import ('),
-        ...     (4, '   date,'),
-        ...     (5, '   datetime,'),
-        ...     (6, ')'),
-        ... ])))
-        >>> assert parsed == [
-        ...     ('from __future__ import unicode_literals', [1]),
-        ...     ('import os.path', [2]),
-        ...     ('from datetime import date,datetime', [3, 4, 5, 6])
-        ... ]
+        Iterator which yields tuple of lines strings composing
+        a single import statement as well as teh line numbers
+        on which the import statement was found.
     """
     while True:
 
@@ -121,47 +126,43 @@ def find_imports_from_lines(iterator):
                 line_numbers.append(line_number)
                 line_imports.append(line)
 
-        # remove unneeded characters
-        line_imports = map(lambda i: CHARS.sub('', i), line_imports)
-
-        # now that extra characters are removed
-        # strip each line
-        line_imports = list_strip(line_imports)
-
-        # handle line breaks which can cause
-        # incorrect syntax when recombined
-        # mostly these are cases when line break
-        # happens around either "import" or "as"
-        # e.g. from long.import.here \n import foo
-        # e.g. import package \n as foo
-        for word in ('import', 'as'):
-            kwargs = {
-                'startswith': {
-                    'word': word,
-                    'pre': '',
-                    'post': ' ',
-                },
-                'endswith': {
-                    'word': word,
-                    'pre': ' ',
-                    'post': '',
-                }
-            }
-            for f, k in kwargs.items():
-                line_imports = list(map(
-                    lambda i: (i if not getattr(i, f)('{pre}{word}{post}'
-                                                      ''.format(**k))
-                               else '{post}{i}{pre}'.format(i=i, **k)),
-                    line_imports
-                ))
-
-        import_line = ''.join(line_imports).strip()
-
-        # strip ending comma if there
-        if import_line.endswith(','):
-            import_line = import_line[:-1]
-
-        yield import_line, line_numbers
+        yield line_imports, line_numbers
+
+
+def tokenize_import_lines(import_lines):
+    tokens = []
+
+    for n, line in enumerate(import_lines):
+        _tokens = []
+        words = filter(None, TOKEN_REGEX.split(line))
+
+        for i, word in enumerate(words):
+            token = Token(word)
+            # tokenize same-line comments before "," to allow to associate
+            # a comment with specific import since pure Python
+            # syntax does not do that because # has to be after ","
+            # hence when tokenizing, comment will be associated
+            # with next import which is not desired
+            if token.is_comment and _tokens and _tokens[max(i - 1, 0)] == ',':
+                _tokens.insert(i - 1, token)
+            else:
+                _tokens.append(token)
+
+        tokens.extend(_tokens)
+
+    # combine tokens between \\ newline escape
+    segments = list_split(tokens, '\\')
+    tokens = [Token('')]
+    for i, segment in enumerate(segments):
+        # don't add to previous token if it is a ","
+        if all((tokens[-1] not in SPECIAL_TOKENS,
+                not i or segment[0] not in SPECIAL_TOKENS)):
+            tokens[-1] += segment[0]
+        else:
+            tokens.append(segment[0])
+        tokens += segment[1:]
+
+    return [Token(i) for i in tokens]
 
 
 def parse_import_statement(stem, line_numbers, **kwargs):
@@ -232,24 +233,56 @@ def parse_statements(iterable, **kwargs):
     statements : generator
         Generator which yields ``ImportStatement`` instances.
     """
-    for import_line, line_numbers in iterable:
-
-        if import_line.startswith('import '):
-            stems = import_line.replace('import ', '').strip().split(',')
-
-            for stem in stems:
-                yield parse_import_statement(stem.strip(),
-                                             line_numbers,
-                                             **kwargs)
+    get_comments = lambda j: filter(lambda i: i.is_comment, j)
+    get_not_comments = lambda j: filter(lambda i: not i.is_comment, j)
+    get_first_not_comment = lambda j: next(
+        iter(filter(lambda i: not i.is_comment, j))
+    )
+
+    for import_lines, line_numbers in iterable:
+        tokens = tokenize_import_lines(import_lines)
+
+        if tokens[0] == 'import':
+            for _tokens in list_split(tokens[1:], ','):
+                stem = ' '.join(get_not_comments(_tokens))
+                comments = get_comments(_tokens)
+                yield parse_import_statement(
+                    stem=stem,
+                    line_numbers=line_numbers,
+                    comments=list(comments),
+                    **kwargs
+                )
 
         else:
-            stem, leafs_string = list_strip(
-                import_line.replace('from ', '').split(' import ')
+            stem_tokens, leafs_tokens = list_split(tokens[1:], 'import')
+            stem = ' '.join(get_not_comments(stem_tokens))
+            list_of_leafs = list(list_split(leafs_tokens, ','))
+            statement_comments = set()
+            leafs = []
+
+            for leaf_list in list_of_leafs:
+                first_non_leaf_index = leaf_list.index(
+                    get_first_not_comment(leaf_list)
+                )
+                leaf_stem = ' '.join(get_not_comments(leaf_list))
+                comments = []
+                all_leaf_comments = filter(lambda i: i.is_comment, leaf_list)
+
+                for comment in all_leaf_comments:
+                    if comment.normalized in STATEMENT_COMMENTS:
+                        statement_comments.add(comment)
+                    else:
+                        comment.is_comment_first = (
+                            leaf_list.index(comment) < first_non_leaf_index
+                        )
+                        comments.append(comment)
+
+                leafs.append(ImportLeaf(leaf_stem, comments=comments))
+
+            yield ImportStatement(
+                line_numbers=line_numbers,
+                stem=stem,
+                leafs=leafs,
+                comments=list(statement_comments),
+                **kwargs
             )
-            leafs = filter(None, list_strip(leafs_string.split(',')))
-            leafs = list(map(ImportLeaf, leafs))
-
-            yield ImportStatement(line_numbers,
-                                  stem,
-                                  leafs,
-                                  **kwargs)
diff --git a/importanize/statements.py b/importanize/statements.py
index 0d7e24e..22eba49 100755
--- a/importanize/statements.py
+++ b/importanize/statements.py
@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 from __future__ import print_function, unicode_literals
+import itertools
 import operator
 import re
 
@@ -23,7 +24,7 @@ class ImportLeaf(ComparatorMixin):
     Also aliased modules are supported (e.g. using ``as``).
     """
 
-    def __init__(self, name):
+    def __init__(self, name, comments=None):
         as_name = None
 
         if ' as ' in name:
@@ -34,6 +35,7 @@ class ImportLeaf(ComparatorMixin):
 
         self.name = name
         self.as_name = as_name
+        self.comments = comments or []
 
     def as_string(self):
         string = self.name
@@ -100,11 +102,13 @@ class ImportStatement(ComparatorMixin):
         List of ``ImportLeaf`` instances
     """
 
-    def __init__(self, line_numbers, stem, leafs=None, **kwargs):
+    def __init__(self, line_numbers, stem, leafs=None,
+                 comments=None, **kwargs):
         self.line_numbers = line_numbers
         self.stem = stem
         self.leafs = leafs or []
-        self.artifacts = kwargs.get('artifacts', {})
+        self.comments = comments or []
+        self.file_artifacts = kwargs.get('file_artifacts', {})
 
     @property
     def unique_leafs(self):
@@ -132,17 +136,58 @@ class ImportStatement(ComparatorMixin):
             )
 
     def formatted(self):
+        leafs = self.unique_leafs
         string = self.as_string()
-        if len(string) > 80 and len(self.leafs) > 1:
-            sep = '{}    '.format(self.artifacts.get('sep', '\n'))
-            string = (
-                'from {} import ({}{},{})'
-                ''.format(self.stem, sep,
-                          (',{}'.format(sep)
-                           .join(map(operator.methodcaller('as_string'),
-                                     self.unique_leafs))),
-                          self.artifacts.get('sep', '\n'))
-            )
+        get_normalized = lambda i: map(
+            operator.attrgetter('normalized'), i
+        )
+
+        all_comments = (
+            self.comments + list(itertools.chain(
+                *list(map(operator.attrgetter('comments'), leafs))
+            ))
+        )
+
+        if len(all_comments) == 1:
+            string += '  # {}'.format(' '.join(get_normalized(all_comments)))
+
+        if any((len(string) > 80 and len(leafs) > 1,
+                len(all_comments) > 1)):
+            sep = '{}    '.format(self.file_artifacts.get('sep', '\n'))
+
+            string = 'from {} import ('.format(self.stem)
+
+            if self.comments:
+                string += '  # {}'.format(' '.join(
+                    get_normalized(self.comments)
+                ))
+
+            for leaf in leafs:
+                string += sep
+
+                first_comments = list(filter(
+                    lambda i: i.is_comment_first,
+                    leaf.comments
+                ))
+                if first_comments:
+                    string += sep.join(
+                        '# {}'.format(i)
+                        for i in get_normalized(first_comments)
+                    ) + sep
+
+                string += '{},'.format(leaf.as_string())
+
+                inline_comments = list(filter(
+                    lambda i: not i.is_comment_first,
+                    leaf.comments
+                ))
+                if inline_comments:
+                    string += '  # {}'.format(
+                        ' '.join(get_normalized(inline_comments))
+                    )
+
+            string += '{})'.format(self.file_artifacts.get('sep', '\n'))
+
         return string
 
     def __hash__(self):
diff --git a/importanize/utils.py b/importanize/utils.py
index 188b123..d81f20e 100644
--- a/importanize/utils.py
+++ b/importanize/utils.py
@@ -55,3 +55,17 @@ def list_strip(data):
 def read(path):
     with open(path, 'rb') as fid:
         return fid.read().decode('utf-8')
+
+
+def list_split(iterable, split):
+    segment = []
+
+    for i in iterable:
+        if i == split:
+            yield segment
+            segment = []
+        else:
+            segment.append(i)
+
+    if segment:
+        yield segment
diff --git a/setup.py b/setup.py
index 7f1c331..4de6350 100755
--- a/setup.py
+++ b/setup.py
@@ -1,8 +1,9 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-from __future__ import unicode_literals, print_function
+from __future__ import print_function, unicode_literals
 import os
-from setuptools import setup, find_packages
+
+from setuptools import find_packages, setup
 
 from importanize import __author__, __version__
 


Test patch:
diff --git a/tests/test_data/normal.txt b/tests/test_data/input.txt
similarity index 56%
rename from tests/test_data/normal.txt
rename to tests/test_data/input.txt
index 553ab4d..8123578 100644
--- a/tests/test_data/normal.txt
+++ b/tests/test_data/input.txt
@@ -12,6 +12,13 @@ from a import b
 from a.b import c
 from a.b import d
 
-
+import something # with comment
+from other.package.subpackage.module.submodule import CONSTANT, Klass, foo, bar, rainbows # noqa
+from other import(
+    something, # something comment
+    something_else, # noqa
+    #lots of happy things below
+    and_rainbows
+)
 
 # stuff here
diff --git a/tests/test_data/normal_expected.txt b/tests/test_data/output.txt
similarity index 54%
rename from tests/test_data/normal_expected.txt
rename to tests/test_data/output.txt
index 63f1f74..e1c57e6 100644
--- a/tests/test_data/normal_expected.txt
+++ b/tests/test_data/output.txt
@@ -4,8 +4,22 @@ from __future__ import print_function, unicode_literals
 import datetime
 from os import path as ospath
 
+import something  # with comment
 from a import b
 from a.b import c, d
+from other import (  # noqa
+    # lots of happy things below
+    and_rainbows,
+    something,  # something comment
+    something_else,
+)
+from other.package.subpackage.module.submodule import (  # noqa
+    CONSTANT,
+    Klass,
+    bar,
+    foo,
+    rainbows,
+)
 from package.subpackage.module.submodule import (
     CONSTANT,
     Klass,
diff --git a/tests/test_groups.py b/tests/test_groups.py
index 7d1ec49..fb1e386 100644
--- a/tests/test_groups.py
+++ b/tests/test_groups.py
@@ -116,7 +116,7 @@ class TestBaseImportGroup(unittest.TestCase):
         )
 
     def test_as_string_with_artifacts(self):
-        group = BaseImportGroup(artifacts={'sep': '\r\n'})
+        group = BaseImportGroup(file_artifacts={'sep': '\r\n'})
         group.statements = [ImportStatement([], 'b'),
                             ImportStatement([], 'a')]
 
@@ -145,12 +145,12 @@ class TestBaseImportGroup(unittest.TestCase):
 
     def test_formatted_with_artifacts(self):
         artifacts = {'sep': '\r\n'}
-        group = BaseImportGroup(artifacts=artifacts)
+        group = BaseImportGroup(file_artifacts=artifacts)
         group.statements = [
             ImportStatement(list(), 'b' * 80, [ImportLeaf('c'),
                                                ImportLeaf('d')],
-                            artifacts=artifacts),
-            ImportStatement([], 'a', artifacts=artifacts)
+                            file_artifacts=artifacts),
+            ImportStatement([], 'a', file_artifacts=artifacts)
         ]
 
         self.assertEqual(
@@ -298,17 +298,17 @@ class TestImportGroups(unittest.TestCase):
     def test_formatted_with_artifacts(self):
         artifacts = {'sep': '\r\n'}
 
-        groups = ImportGroups(artifacts=artifacts)
+        groups = ImportGroups(file_artifacts=artifacts)
         groups.groups = [
-            RemainderGroup(artifacts=artifacts),
-            LocalGroup(artifacts=artifacts),
+            RemainderGroup(file_artifacts=artifacts),
+            LocalGroup(file_artifacts=artifacts),
         ]
 
         groups.add_statement_to_group(
-            ImportStatement([], '.a', artifacts=artifacts)
+            ImportStatement([], '.a', file_artifacts=artifacts)
         )
         groups.add_statement_to_group(
-            ImportStatement([], 'foo', artifacts=artifacts)
+            ImportStatement([], 'foo', file_artifacts=artifacts)
         )
 
         self.assertEqual(
diff --git a/tests/test_main.py b/tests/test_main.py
index c4992c3..d39935a 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -1,11 +1,20 @@
 # -*- coding: utf-8 -*-
 from __future__ import print_function, unicode_literals
+import logging
 import os
 import unittest
 
 import mock
-
-from importanize.main import PEP8_CONFIG, run, run_importanize
+import six
+
+from importanize.main import (
+    PEP8_CONFIG,
+    find_config,
+    main,
+    parser,
+    run,
+    run_importanize,
+)
 from importanize.utils import read
 
 
@@ -27,11 +36,15 @@ class TestMain(unittest.TestCase):
     def test_run_importanize_print(self, mock_print):
         test_file = os.path.join(os.path.dirname(__file__),
                                  'test_data',
-                                 'normal.txt')
+                                 'input.txt')
         expected_file = os.path.join(os.path.dirname(__file__),
                                      'test_data',
-                                     'normal_expected.txt')
-        expected = read(expected_file).encode('utf-8')
+                                     'output.txt')
+        expected = (
+            read(expected_file)
+            if six.PY3
+            else read(expected_file).encode('utf-8')
+        )
 
         self.assertTrue(
             run_importanize(test_file,
@@ -44,10 +57,10 @@ class TestMain(unittest.TestCase):
     def test_run_importanize_write(self, mock_open):
         test_file = os.path.join(os.path.dirname(__file__),
                                  'test_data',
-                                 'normal.txt')
+                                 'input.txt')
         expected_file = os.path.join(os.path.dirname(__file__),
                                      'test_data',
-                                     'normal_expected.txt')
+                                     'output.txt')
         expected = read(expected_file).encode('utf-8')
 
         self.assertTrue(
@@ -157,3 +170,81 @@ class TestMain(unittest.TestCase):
             conf,
         )
         mock_parser.error.assert_called_once_with(mock.ANY)
+
+    @mock.patch('os.path.exists')
+    @mock.patch('os.getcwd')
+    def test_find_config(self, mock_getcwd, mock_exists):
+        mock_getcwd.return_value = '/path/to/importanize'
+        mock_exists.side_effect = False, True
+
+        config, found = find_config()
+
+        self.assertEqual(config, '/path/to/.importanizerc')
+        self.assertTrue(bool(found))
+
+    @mock.patch(TESTING_MODULE + '.run')
+    @mock.patch('logging.getLogger')
+    @mock.patch.object(parser, 'parse_args')
+    def test_main_without_config(self,
+                                 mock_parse_args,
+                                 mock_get_logger,
+                                 mock_run):
+        args = mock.MagicMock(
+            verbose=1,
+            version=False,
+            path=['/path/../'],
+            config=None,
+        )
+        mock_parse_args.return_value = args
+
+        main()
+
+        mock_parse_args.assert_called_once_with()
+        mock_get_logger.assert_called_once_with('')
+        mock_get_logger().setLevel.assert_called_once_with(logging.INFO)
+        mock_run.assert_called_once_with('/', PEP8_CONFIG, args)
+
+    @mock.patch(TESTING_MODULE + '.read')
+    @mock.patch(TESTING_MODULE + '.run')
+    @mock.patch('json.loads')
+    @mock.patch('logging.getLogger')
+    @mock.patch.object(parser, 'parse_args')
+    def test_main_with_config(self,
+                              mock_parse_args,
+                              mock_get_logger,
+                              mock_loads,
+                              mock_run,
+                              mock_read):
+        args = mock.MagicMock(
+            verbose=1,
+            version=False,
+            path=['/path/../'],
+            config=mock.sentinel.config,
+        )
+        mock_parse_args.return_value = args
+
+        main()
+
+        mock_parse_args.assert_called_once_with()
+        mock_get_logger.assert_called_once_with('')
+        mock_get_logger().setLevel.assert_called_once_with(logging.INFO)
+        mock_read.assert_called_once_with(mock.sentinel.config)
+        mock_loads.assert_called_once_with(mock_read.return_value)
+        mock_run.assert_called_once_with('/', mock_loads.return_value, args)
+
+    @mock.patch(TESTING_MODULE + '.print', create=True)
+    @mock.patch('sys.exit')
+    @mock.patch.object(parser, 'parse_args')
+    def test_main_version(self, mock_parse_args, mock_exit, mock_print):
+        mock_exit.side_effect = SystemExit
+        mock_parse_args.return_value = mock.MagicMock(
+            verbose=1,
+            version=True,
+        )
+
+        with self.assertRaises(SystemExit):
+            main()
+
+        mock_parse_args.assert_called_once_with()
+        mock_exit.assert_called_once_with(0)
+        mock_print.assert_called_once_with(mock.ANY)
diff --git a/tests/test_parser.py b/tests/test_parser.py
index 3e6d828..3c8a22a 100644
--- a/tests/test_parser.py
+++ b/tests/test_parser.py
@@ -6,33 +6,47 @@ import mock
 import six
 
 from importanize.parser import (
+    Token,
     find_imports_from_lines,
-    get_artifacts,
+    get_file_artifacts,
     parse_statements,
+    tokenize_import_lines,
 )
 
 
 TESTING_MODULE = 'importanize.parser'
 
 
+class TestToken(unittest.TestCase):
+    def test_is_comment(self):
+        self.assertTrue(Token('# noqa').is_comment)
+        self.assertFalse(Token('foo').is_comment)
+
+    def test_normalized(self):
+        self.assertEqual(Token('foo').normalized, 'foo')
+        self.assertEqual(Token('#noqa').normalized, 'noqa')
+        self.assertEqual(Token('# noqa').normalized, 'noqa')
+        self.assertEqual(Token('#  noqa').normalized, ' noqa')
+
+
 class TestParsing(unittest.TestCase):
     @mock.patch(TESTING_MODULE + '.read')
-    def test_get_artifacts(self, mock_read):
+    def test_get_file_artifacts(self, mock_read):
         mock_read.return_value = 'Hello\nWorld\n'
-        actual = get_artifacts(mock.sentinel.path)
+        actual = get_file_artifacts(mock.sentinel.path)
         self.assertDictEqual(actual, {
             'sep': '\n',
         })
         mock_read.assert_called_once_with(mock.sentinel.path)
 
         mock_read.return_value = 'Hello\r\nWorld\n'
-        actual = get_artifacts(mock.sentinel.path)
+        actual = get_file_artifacts(mock.sentinel.path)
         self.assertDictEqual(actual, {
             'sep': '\r\n',
         })
 
         mock_read.return_value = 'Hello'
-        actual = get_artifacts(mock.sentinel.path)
+        actual = get_file_artifacts(mock.sentinel.path)
         self.assertDictEqual(actual, {
             'sep': '\n',
         })
@@ -46,68 +60,69 @@ class TestParsing(unittest.TestCase):
     def test_parsing(self):
         self._test_import_parsing(
             ('import a',),
-            ('import a', [0]),
+            (['import a'], [0]),
         )
         self._test_import_parsing(
             ('import a, b',),
-            ('import a, b', [0]),
+            (['import a, b'], [0]),
         )
         self._test_import_parsing(
             ('import a, b as c',),
-            ('import a, b as c', [0]),
+            (['import a, b as c'], [0]),
         )
         self._test_import_parsing(
             ('from a import b',),
-            ('from a import b', [0]),
+            (['from a import b'], [0]),
         )
         self._test_import_parsing(
             ('from a.b import c',),
-            ('from a.b import c', [0]),
+            (['from a.b import c'], [0]),
         )
         self._test_import_parsing(
             ('from a.b import c,\\',
              '    d'),
-            ('from a.b import c,d', [0, 1]),
+            (['from a.b import c,\\',
+              '    d'], [0, 1]),
         )
         self._test_import_parsing(
             ('from a.b import c, \\',
              '    d'),
-            ('from a.b import c,d', [0, 1]),
+            (['from a.b import c, \\',
+              '    d'], [0, 1]),
         )
         self._test_import_parsing(
             ('from a.b \\',
              '    import c'),
-            ('from a.b import c', [0, 1]),
+            (['from a.b \\',
+              '    import c'], [0, 1]),
         )
         self._test_import_parsing(
             ('import a.b \\',
              '    as c'),
-            ('import a.b as c', [0, 1]),
+            (['import a.b \\',
+              '    as c'], [0, 1]),
         )
         self._test_import_parsing(
             ('from a.b import \\',
              '    c,\\',
              '    d,'),
-            ('from a.b import c,d', [0, 1, 2]),
+            (['from a.b import \\',
+              '    c,\\',
+              '    d,'], [0, 1, 2]),
         )
         self._test_import_parsing(
             ('from a.b import (c,',
              '    d)'),
-            ('from a.b import c,d', [0, 1]),
+            (['from a.b import (c,',
+              '    d)'], [0, 1]),
         )
         self._test_import_parsing(
             ('from a.b import (c, ',
              '    d',
              ')'),
-            ('from a.b import c,d', [0, 1, 2]),
-        )
-        self._test_import_parsing(
-            ('from a.b import (',
-             '    c,',
-             '    d,',
-             ')',
-             'foo'),
-            ('from a.b import c,d', [0, 1, 2, 3]),
+            (['from a.b import (c, ',
+              '    d',
+              ')'], [0, 1, 2]),
         )
         self._test_import_parsing(
             ('"""',
@@ -118,7 +133,10 @@ class TestParsing(unittest.TestCase):
              '    d,',
              ')',
              'foo'),
-            ('from a.b import c,d', [3, 4, 5, 6]),
+            (['from a.b import (',
+              '    c,',
+              '    d,',
+              ')'], [3, 4, 5, 6]),
         )
         self._test_import_parsing(
             ('""" from this shall not import """',
@@ -127,7 +145,10 @@ class TestParsing(unittest.TestCase):
              '    d,',
              ')',
              'foo'),
-            ('from a.b import c,d', [1, 2, 3, 4]),
+            (['from a.b import (',
+              '    c,',
+              '    d,',
+              ')'], [1, 2, 3, 4]),
         )
         self._test_import_parsing(
             ('"""',
@@ -145,10 +166,89 @@ class TestParsing(unittest.TestCase):
             tuple(),
         )
 
+    def test_tokenize_import_lines(self):
+        data = (
+            (
+                ['import a.\\',
+                 '    b'],
+                ['import',
+                 'a.b']
+            ),
+            (
+                ['from __future__ import unicode_literals,\\',
+                 '    print_function'],
+                ['from',
+                 '__future__',
+                 'import',
+                 'unicode_literals',
+                 ',',
+                 'print_function']
+            ),
+            (
+                ['from a import b,  # noqa',
+                 '    c'],
+                ['from',
+                 'a',
+                 'import',
+                 'b',
+                 '# noqa',
+                 ',',
+                 'c']
+            ),
+            (
+                ['import a,\\',
+                 '    b'],
+                ['import',
+                 'a',
+                 ',',
+                 'b']
+            ),
+            (
+                ['from a import\\',
+                 '    b'],
+                ['from',
+                 'a',
+                 'import',
+                 'b']
+            ),
+            (
+                ['from a\\',
+                 '    import b'],
+                ['from',
+                 'a',
+                 'import',
+                 'b']
+            ),
+            (
+                ['import a\\',
+                 '    as b'],
+                ['import',
+                 'a',
+                 'as',
+                 'b']
+            ),
+            (
+                ['from something import foo, bar  # noqa'],
+                ['from',
+                 'something',
+                 'import',
+                 'foo',
+                 ',',
+                 'bar',
+                 '# noqa']
+            ),
+        )
+        for _data, expected in data:
+            self.assertListEqual(
+                tokenize_import_lines(iter(_data)),
+                expected
+            )
+
     def _test_import_string_matches(self, string, expected):
+        data = string if isinstance(string, (list, tuple)) else [string]
         self.assertEqual(
             six.text_type(
-                list(parse_statements([(string, [1])]))[0]
+                next(parse_statements([(data, [1])]))
             ),
             expected
         )
@@ -238,3 +338,10 @@ class TestParsing(unittest.TestCase):
             'from a import b as e,d as g, c as c',
             'from a import b as e, c, d as g',
         )
+        self._test_import_string_matches(
+            ['from a import (',
+             '    b as e,  # foo',
+             '    d as g  # noqa',
+             ')'],
+            'from a import b as e, d as g',
+        )
diff --git a/tests/test_statements.py b/tests/test_statements.py
index 7d740fd..289f5bb 100644
--- a/tests/test_statements.py
+++ b/tests/test_statements.py
@@ -5,6 +5,7 @@ import unittest
 import mock
 import six
 
+from importanize.parser import Token
 from importanize.statements import ImportLeaf, ImportStatement
 
 
@@ -147,11 +148,15 @@ class TestImportStatement(unittest.TestCase):
         _test('a.b', ['e', 'c as d', 'e'], 'from a.b import c as d, e')
 
     def test_formatted(self):
-        def _test(stem, leafs, expected, sep='\n', **kwargs):
+        def _test(stem, leafs, expected, sep='\n', comments=None, **kwargs):
             statement = ImportStatement(
                 list(),
                 stem,
-                list(map(ImportLeaf, leafs)),
+                list(map((lambda i:
+                          i if isinstance(i, ImportLeaf)
+                          else ImportLeaf(i)),
+                         leafs)),
+                comments=comments,
                 **kwargs
             )
             self.assertEqual(statement.formatted(),
@@ -171,8 +176,29 @@ class TestImportStatement(unittest.TestCase):
                '    {},'.format('b' * 20),
                '    {},'.format('c' * 20),
                ')'],
-              '\r\n',
-              artifacts={'sep': '\r\n'})
+              sep='\r\n',
+              file_artifacts={'sep': '\r\n'})
+        _test('foo', [],
+              ['import foo  # comment'],
+              comments=[Token('# comment')])
+        _test('foo', [ImportLeaf('bar', comments=[Token('#comment')])],
+              ['from foo import bar  # comment'])
+        _test('something', [ImportLeaf('foo'),
+                            ImportLeaf('bar')],
+              ['from something import bar, foo  # noqa'],
+              comments=[Token('# noqa')])
+        _test('foo',
+              [ImportLeaf('bar', comments=[Token('#hello')]),
+               ImportLeaf('rainbows', comments=[Token('#world')]),
+               ImportLeaf('zzz', comments=[Token('#and lots of sleep',
+                                                 is_comment_first=True)])],
+              ['from foo import (  # noqa',
+               '    bar,  # hello',
+               '    rainbows,  # world',
+               '    # and lots of sleep',
+               '    zzz,',
+               ')'],
+              comments=[Token('#noqa')])
 
     def test_str(self):
         statement = ImportStatement([], 'a')