Add support for recursive pycache removal.

Author: dave-shawley

HISTORY

@@ -4,6 +4,7 @@ Changelog
 * 0.0.3
 
   - Support removal of virtual environment directories.
+  - Support recursive removal of *__pycache__* directories.
 
 * 0.0.2 (17-Nov-2014)
 

README.rst

@@ -75,6 +75,9 @@ few new command line parameters.
    ``$VIRTUAL_ENV`` environment variable.  The name of the directory can
    also be specified using the ``--virtualenv-dir`` command line option.
 
+``setup.py clean --pycache``
+   Recursively removes directories named *__pycache__*.
+
 Where can I get this extension from?
 ------------------------------------
 +---------------+-----------------------------------------------------+

setupext/janitor.py

@@ -42,6 +42,7 @@ def initialize_options(self):
         self.eggs = False
         self.egg_base = None
         self.environment = False
+        self.pycache = False
         self.virtualenv_dir = None
 
     def finalize_options(self):
@@ -81,6 +82,11 @@ def run(self):
         if self.environment and self.virtualenv_dir:
             dir_names.add(self.virtualenv_dir)
 
+        if self.pycache:
+            for root, dirs, _ in os.walk(os.curdir):
+                if '__pycache__' in dirs:
+                    dir_names.add(os.path.join(root, '__pycache__'))
+
         for dir_name in dir_names:
             if os.path.exists(dir_name):
                 dir_util.remove_tree(dir_name, dry_run=self.dry_run)
@@ -110,6 +116,7 @@ def _set_options():
         ('dist', 'd', 'remove distribution directory'),
         ('eggs', None, 'remove egg and egg-info directories'),
         ('environment', 'E', 'remove virtual environment directory'),
+        ('pycache', 'p', 'remove __pycache__ directories'),
 
         ('egg-base=', 'e',
          'directory containing .egg-info directories '
@@ -119,6 +126,7 @@ def _set_options():
          '(default: value of VIRTUAL_ENV environment variable)'),
     ])
     CleanCommand.boolean_options = _CleanCommand.boolean_options[:]
-    CleanCommand.boolean_options.extend(['dist', 'eggs', 'environment'])
+    CleanCommand.boolean_options.extend(
+        ['dist', 'eggs', 'environment', 'pycache'])
 
 _set_options()

tests.py

@@ -184,6 +184,7 @@ def test_that_directories_are_not_removed_in_dry_run_mode(self):
 
 
 class VirtualEnvironmentCleanupTests(DirectoryCleanupMixin, unittest.TestCase):
+
     def setUp(self):
         super(VirtualEnvironmentCleanupTests, self).setUp()
         self._saved_venv = os.environ.get('VIRTUAL_ENV', None)
@@ -216,3 +217,41 @@ def test_that_janitor_does_not_fail_when_envdir_is_missing(self):
     def test_that_janitor_does_not_fail_when_no_dir_specified(self):
         os.environ.pop('VIRTUAL_ENV', None)
         run_setup('clean', '--environment')
+
+
+class PycacheCleanupTests(DirectoryCleanupMixin, unittest.TestCase):
+
+    def setUp(self):
+        super(PycacheCleanupTests, self).setUp()
+        self.test_root = self.create_directory('test-root')
+        starting_dir = os.curdir
+        self.addCleanup(os.chdir, starting_dir)
+        os.chdir(self.test_root)
+
+    def mkdirs(self, *dirs):
+        for d in dirs:
+            os.makedirs(d)
+        return dirs[:]
+
+    def test_that_pycache_directories_are_removed(self):
+        all_dirs = self.mkdirs(
+            os.path.join(self.test_root, '__pycache__'),
+            os.path.join(self.test_root, 'a', '__pycache__'),
+            os.path.join(self.test_root, 'a', 'b', '__pycache__'),
+            os.path.join(self.test_root, 'b', '__pycache__'),
+        )
+
+        run_setup('clean', '--pycache')
+        for cache_dir in all_dirs:
+            self.assert_path_does_not_exist(cache_dir)
+
+    def test_that_janitor_does_not_fail_when_cache_parent_is_removed(self):
+        all_dirs = self.mkdirs(
+            os.path.join(self.test_root, 'dist'),
+            os.path.join(self.test_root, 'dist', '__pycache__', '__pycache__'),
+            os.path.join(self.test_root, 'dist', 'foo', '__pycache__'),
+        )
+
+        run_setup('clean', '--dist', '--pycache')
+        for cache_dir in all_dirs:
+            self.assert_path_does_not_exist(cache_dir)