diff --git a/README.rst b/README.rst index 7516bd4..4ecd99a 100644 --- a/README.rst +++ b/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 diff --git a/autoflake.py b/autoflake.py index b192e03..adb3a56 100755 --- a/autoflake.py +++ b/autoflake.py @@ -117,6 +117,20 @@ 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""" + pattern = r':\ (.+)$' + 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 +286,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_single_star_import=False, remove_all_unused_imports=False, remove_unused_variables=False): """Yield code with unused imports removed.""" @@ -289,6 +304,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_single_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 len(undefined_names) == 0: + 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 +340,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 +484,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_single_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 +501,7 @@ def fix_code(source, additional_imports=None, remove_all_unused_imports=False, filter_code( source, additional_imports=additional_imports, + expand_single_star_import=expand_single_star_import, remove_all_unused_imports=remove_all_unused_imports, remove_unused_variables=remove_unused_variables)))) @@ -486,6 +523,7 @@ def fix_file(filename, args, standard_out): filtered_source = fix_code( source, additional_imports=args.imports.split(',') if args.imports else None, + expand_single_star_import=args.expand_single_star_import, remove_all_unused_imports=args.remove_all_unused_imports, remove_unused_variables=args.remove_unused_variables) @@ -569,9 +607,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-single-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', diff --git a/test_autoflake.py b/test_autoflake.py index d897e5e..763a740 100755 --- a/test_autoflake.py +++ b/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_single_star_import(self): + self.assertEqual( + """\ +from math import sin +sin(1) +""", + ''.join(autoflake.filter_code("""\ +from math import * +sin(1) +""", expand_single_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_single_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_single_star_import=True))) + + def test_multiline_import(self): self.assertTrue(autoflake.multiline_import(r"""\ import os, \ diff --git a/test_fuzz.py b/test_fuzz.py index ac7de62..b767c79 100755 --- a/test_fuzz.py +++ b/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-single-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_single_star_import: + options.append('--expand-single-star-import') + if args.imports: options.append('--imports=' + args.imports)