Remove imports and variables

pyformat.py

@@ -41,10 +41,14 @@
 __version__ = '0.6'
 
 
-def formatters(aggressive, apply_config, filename=''):
+def formatters(aggressive, apply_config, filename='',
+               remove_all_unused_imports=False, remove_unused_variables=False):
     """Return list of code formatters."""
     if aggressive:
-        yield autoflake.fix_code
+        yield lambda code: autoflake.fix_code(
+            code,
+            remove_all_unused_imports=remove_all_unused_imports,
+            remove_unused_variables=remove_unused_variables)
         autopep8_options = autopep8.parse_args(
             [filename] + int(aggressive) * ['--aggressive'],
             apply_config=apply_config)
@@ -57,11 +61,15 @@ def formatters(aggressive, apply_config, filename=''):
     yield unify.format_code
 
 
-def format_code(source, aggressive=False, apply_config=False, filename=''):
+def format_code(source, aggressive=False, apply_config=False, filename='',
+                remove_all_unused_imports=False,
+                remove_unused_variables=False):
     """Return formatted source code."""
     formatted_source = source
 
-    for fix in formatters(aggressive, apply_config, filename):
+    for fix in formatters(
+            aggressive, apply_config, filename,
+            remove_all_unused_imports, remove_unused_variables):
         formatted_source = fix(formatted_source)
 
     return formatted_source
@@ -81,10 +89,13 @@ def format_file(filename, args, standard_out):
     if not source:
         return False
 
-    formatted_source = format_code(source,
-                                   aggressive=args.aggressive,
-                                   apply_config=args.config,
-                                   filename=filename)
+    formatted_source = format_code(
+        source,
+        aggressive=args.aggressive,
+        apply_config=args.config,
+        filename=filename,
+        remove_all_unused_imports=args.remove_all_unused_imports,
+        remove_unused_variables=args.remove_unused_variables)
 
     if source != formatted_source:
         if args.in_place:
@@ -160,6 +171,12 @@ def parse_args(argv):
                         help='drill down directories recursively')
     parser.add_argument('-a', '--aggressive', action='count', default=0,
                         help='use more aggressive formatters')
+    parser.add_argument('--remove-all-unused-imports', action='store_true',
+                        help='remove all unused imports, '
+                             'not just standard library '
+                             '(requires "aggressive")')
+    parser.add_argument('--remove-unused-variables', action='store_true',
+                        help='remove unused variables (requires "aggressive")')
     parser.add_argument('-j', '--jobs', type=int, metavar='n', default=1,
                         help='number of parallel jobs; '
                              'match CPU count if value is less than 1')
@@ -185,6 +202,12 @@ def parse_args(argv):
         import multiprocessing
         args.jobs = multiprocessing.cpu_count()
 
+    if args.remove_all_unused_imports and not args.aggressive:
+        parser.error('--remove-all-unused-imports requires --aggressive')
+
+    if args.remove_unused_variables and not args.aggressive:
+        parser.error('--remove-unused-variables requires --aggressive')
+
     return args
 
 

test_pyformat.py

@@ -60,6 +60,18 @@ def test_format_code_with_unicode(self):
                          pyformat.format_code(
                              'x = "abc" \\\n"ö"\n'))
 
+    def test_format_code_with_remove_all_unused_imports(self):
+        self.assertEqual("x = 'abc' \\\n    'ö'\n",
+                         pyformat.format_code(
+                             'import mymodule\nx = "abc" \\\n"ö"\n', aggressive=True,
+                             remove_all_unused_imports=True))
+
+    def test_format_code_with_remove_unused_variables(self):
+        self.assertEqual("def test():\n    return 42\n",
+                         pyformat.format_code(
+                             'def test():\n    x = 4\n    return 42', aggressive=True,
+                             remove_unused_variables=True))
+
     def test_format_multiple_files(self):
         with temporary_file('''\
 if True:
@@ -242,6 +254,22 @@ def test_jobs_less_than_one_should_default_to_cpu_count(self):
 
         self.assertGreater(args.jobs, 0)
 
+    def test_remove_all_unused_imports_requires_aggressive(self):
+        _stderr = sys.stderr
+        sys.stderr = io.StringIO()
+        self.assertRaises(
+            SystemExit, pyformat.parse_args,
+            ['my_fake_program', '--remove-all-unused-imports', __file__])
+        sys.stderr = _stderr
+
+    def test_remove_unused_variables_requires_aggressive(self):
+        _stderr = sys.stderr
+        sys.stderr = io.StringIO()
+        self.assertRaises(
+            SystemExit, pyformat.parse_args,
+            ['my_fake_program', '--remove-unused-variables', __file__])
+        sys.stderr = _stderr
+
     def test_ignore_hidden_directories(self):
         with temporary_directory() as directory:
             with temporary_directory(prefix='.',
@@ -303,6 +331,48 @@ def test_exclude(self):
                     '',
                     output_file.getvalue().strip())
 
+    def test_remove_all_unused_imports(self):
+        with temporary_file("""\
+import mymodule
+
+def test():
+    return 42
+""") as filename:
+            output_file = io.StringIO()
+            pyformat._main(argv=['my_fake_program',
+                                 '--in-place',
+                                 '--aggressive',
+                                 '--remove-all-unused-imports',
+                                 filename],
+                           standard_out=output_file,
+                           standard_error=None)
+            with open(filename) as f:
+                self.assertEqual('''\
+
+def test():
+    return 42
+''', f.read())
+
+    def test_remove_unused_variables(self):
+        with temporary_file("""\
+def test():
+    x = 43
+    return 42
+""") as filename:
+            output_file = io.StringIO()
+            pyformat._main(argv=['my_fake_program',
+                                 '--in-place',
+                                 '--aggressive',
+                                 '--remove-unused-variables',
+                                 filename],
+                           standard_out=output_file,
+                           standard_error=None)
+            with open(filename) as f:
+                self.assertEqual('''\
+def test():
+    return 42
+''', f.read())
+
     def test_end_to_end(self):
         with temporary_file("""\
 import os