Add option to expand single star import

README.rst

@@ -91,11 +91,11 @@ To remove unused variables, use the ``--remove-unused-variables`` option.
 Below is the full listing of options::
 
     usage: autoflake [-h] [-i] [-r] [--imports IMPORTS]
-                     [--remove-all-unused-imports] [--remove-unused-variables]
-                     [--version]
+                     [--expand-single-star-import] [--remove-all-unused-imports]
+                     [--remove-unused-variables] [--version]
                      files [files ...]
 
-    Removes unused imports as reported by pyflakes.
+    Removes unused imports and unused variables as reported by pyflakes.
 
     positional arguments:
       files                 files to format
@@ -107,9 +107,11 @@ Below is the full listing of options::
       --imports IMPORTS     by default, only unused standard library imports are
                             removed; specify a comma-separated list of additional
                             modules/packages
+      --expand-single-star-import
+                            expand wildcard star import with undefined names
       --remove-all-unused-imports
                             remove all unused imports (not just those from the
-                            standard library
+                            standard library)
       --remove-unused-variables
                             remove unused variables
       --version             show program's version number and exit

autoflake.py

@@ -117,6 +117,19 @@ def unused_import_module_name(messages):
             if module_name:
                 yield (message.lineno, module_name)
 
+def star_import_used_line_numbers(messages):
+    """Yield line number of star import usage"""
+    for message in messages:
+        if isinstance(message, pyflakes.messages.ImportStarUsed):
+            yield message.lineno
+
+def star_import_usage_undefined_name(messages):
+    """Yield line number, undefined name, and its possible origin module"""
+    for message in messages:
+        if isinstance(message, pyflakes.messages.ImportStarUsage):
+            undefined_name = message.message_args[0]
+            module_name = message.message_args[1]
+            yield (message.lineno, undefined_name, module_name)
 
 def unused_variable_line_numbers(messages):
     """Yield line numbers of unused variables."""
@@ -272,7 +285,8 @@ def break_up_import(line):
                     for i in sorted(imports.split(','))])
 
 
-def filter_code(source, additional_imports=None,
+def filter_code(source, additional_imports=None, 
+                expand_star_import=False,
                 remove_all_unused_imports=False,
                 remove_unused_variables=False):
     """Yield code with unused imports removed."""
@@ -289,6 +303,22 @@ def filter_code(source, additional_imports=None,
     for line_number, module_name in unused_import_module_name(messages):
         marked_unused_module[line_number].append(module_name)
 
+    if expand_star_import:
+        marked_star_import_line_numbers = frozenset(
+            star_import_used_line_numbers(messages))        
+        if len(marked_star_import_line_numbers) > 1:
+            # Auto expanding only possible for single star import
+            marked_star_import_line_numbers = frozenset()
+        else:
+            undefined_names = []
+            for line_number, undefined_name, _ \
+                 in star_import_usage_undefined_name(messages):
+                undefined_names.append(undefined_name)
+            if not undefined_names:
+                marked_star_import_line_numbers = frozenset()
+    else:
+        marked_star_import_line_numbers = frozenset()
+
     if remove_unused_variables:
         marked_variable_line_numbers = frozenset(
             unused_variable_line_numbers(messages))
@@ -309,11 +339,16 @@ def filter_code(source, additional_imports=None,
                 previous_line=previous_line)
         elif line_number in marked_variable_line_numbers:
             yield filter_unused_variable(line)
+        elif line_number in marked_star_import_line_numbers:
+            yield filter_star_import(line, undefined_names)
         else:
             yield line
 
         previous_line = line
 
+def filter_star_import(line, marked_star_import_undefined_name):
+    undefined_name = sorted(set(marked_star_import_undefined_name))
+    return re.sub(r'\*', ', '.join(undefined_name), line)
 
 def filter_unused_import(line, unused_module, remove_all_unused_imports,
                          imports, previous_line=''):
@@ -448,8 +483,8 @@ def get_line_ending(line):
         return line[non_whitespace_index:]
 
 
-def fix_code(source, additional_imports=None, remove_all_unused_imports=False,
-             remove_unused_variables=False):
+def fix_code(source, additional_imports=None, expand_star_import=False,
+             remove_all_unused_imports=False, remove_unused_variables=False):
     """Return code with all filtering run on it."""
     if not source:
         return source
@@ -465,6 +500,7 @@ def fix_code(source, additional_imports=None, remove_all_unused_imports=False,
                 filter_code(
                     source,
                     additional_imports=additional_imports,
+                    expand_star_import=expand_star_import,
                     remove_all_unused_imports=remove_all_unused_imports,
                     remove_unused_variables=remove_unused_variables))))
 
@@ -486,6 +522,7 @@ def fix_file(filename, args, standard_out):
     filtered_source = fix_code(
         source,
         additional_imports=args.imports.split(',') if args.imports else None,
+        expand_star_import=args.expand_star_import,
         remove_all_unused_imports=args.remove_all_unused_imports,
         remove_unused_variables=args.remove_unused_variables)
 
@@ -569,9 +606,11 @@ def _main(argv, standard_out, standard_error):
                         help='by default, only unused standard library '
                              'imports are removed; specify a comma-separated '
                              'list of additional modules/packages')
+    parser.add_argument('--expand-star-import', action='store_true',
+                        help='expand wildcard star import with undefined names')
     parser.add_argument('--remove-all-unused-imports', action='store_true',
                         help='remove all unused imports (not just those from '
-                             'the standard library')
+                             'the standard library)')
     parser.add_argument('--remove-unused-variables', action='store_true',
                         help='remove unused variables')
     parser.add_argument('--version', action='version',

test_autoflake.py

@@ -96,6 +96,17 @@ def test_get_indentation(self):
         self.assertEqual(' \t ', autoflake.get_indentation(' \t abc  \n\t'))
         self.assertEqual('', autoflake.get_indentation('    '))
 
+    def test_filter_star_import(self):
+        self.assertEqual(
+            'from math import cos',
+            autoflake.filter_star_import('from math import *',
+                                         ['cos']))
+
+        self.assertEqual(
+            'from math import cos, sin',
+            autoflake.filter_star_import('from math import *',
+                                         ['sin', 'cos']))
+
     def test_filter_unused_variable(self):
         self.assertEqual('foo()',
                          autoflake.filter_unused_variable('x = foo()'))
@@ -292,6 +303,45 @@ def test_filter_code_should_respect_noqa(self):
 x = 1
 """)))
 
+    def test_filter_code_expand_star_import(self):
+        self.assertEqual(
+            """\
+from math import sin
+sin(1)
+""",
+            ''.join(autoflake.filter_code("""\
+from math import *
+sin(1)
+""", expand_star_import=True)))
+
+        self.assertEqual(
+            """\
+from math import cos, sin
+sin(1)
+cos(1)
+""",
+            ''.join(autoflake.filter_code("""\
+from math import *
+sin(1)
+cos(1)
+""", expand_star_import=True)))
+
+    def test_filter_code_ignore_multiple_star_import(self):
+        self.assertEqual(
+            """\
+from math import *
+from re import *
+sin(1)
+cos(1)
+""",
+            ''.join(autoflake.filter_code("""\
+from math import *
+from re import *
+sin(1)
+cos(1)
+""", expand_star_import=True)))
+
+
     def test_multiline_import(self):
         self.assertTrue(autoflake.multiline_import(r"""\
 import os, \

test_fuzz.py

@@ -142,6 +142,9 @@ def process_args():
     parser.add_argument('--command', default=AUTOFLAKE_BIN,
                         help='autoflake command (default: %(default)s)')
 
+    parser.add_argument('--expand-star-import', action='store_true',
+                        help='expand wildcard star import with undefined names')
+
     parser.add_argument('--imports',
                         help='pass to the autoflake "--imports" option')
 
@@ -174,6 +177,9 @@ def check(args):
                      if os.path.isdir(path)]
 
     options = []
+    if args.expand_star_import:
+        options.append('--expand-star-import')
+
     if args.imports:
         options.append('--imports=' + args.imports)