Merge pull request #4 from dave-shawley/finish-implementation

Finish implementation

Author: dave-shawley

HISTORY

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

README.rst

@@ -75,6 +75,13 @@ 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__*.
+
+``setup.py clean --all``
+   Remove all of by-products.  This is the same as using ``--dist --egg
+   --environment --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):
@@ -55,6 +56,10 @@ def finalize_options(self):
         if self.egg_base is None:
             self.egg_base = os.curdir
 
+        if self.all:
+            for flag in self.boolean_options:
+                setattr(self, flag, True)
+
         if self.environment and self.virtualenv_dir is None:
             self.virtualenv_dir = os.environ.get('VIRTUAL_ENV', None)
 
@@ -81,6 +86,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 +120,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 +130,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

@@ -84,6 +84,12 @@ def setUpClass(cls):
     def create_directory(cls, dir_name):
         return tempfile.mkdtemp(dir=cls.temp_dir, prefix=dir_name)
 
+    @staticmethod
+    def mkdirs(*dirs):
+        for d in dirs:
+            os.makedirs(d)
+        return dirs[:]
+
     def assert_path_does_not_exist(self, *trailing_segments):
         full_path = os.path.join(*trailing_segments)
         if os.path.exists(full_path):
@@ -184,6 +190,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 +223,83 @@ 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 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)
+
+
+class RemoveAllTests(DirectoryCleanupMixin, unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        super(RemoveAllTests, cls).setUpClass()
+        test_root = cls.create_directory('test-root')
+        starting_dir = os.curdir
+        saved_env_dir = os.environ.get('VIRTUAL_ENV', None)
+
+        def restore():
+            os.chdir(starting_dir)
+            os.environ.pop('VIRTUAL_ENV', None)
+            if saved_env_dir is not None:
+                os.environ['VIRTUAL_ENV'] = saved_env_dir
+
+        os.chdir(test_root)
+        try:
+            all_dirs = cls.mkdirs(
+                '__pycache__',
+                'dist',
+                'foo.egg-info',
+                'env',
+            )
+            cls.pycache_dir = all_dirs[0]
+            cls.dist_dir = all_dirs[1]
+            cls.egg_dir = all_dirs[2]
+            cls.env_dir = all_dirs[3]
+            os.environ['VIRTUAL_ENV'] = cls.env_dir
+            run_setup('clean', '--all')
+            restore()
+        except:
+            restore()
+            raise
+
+    def test_that_pycache_is_removed(self):
+        self.assert_path_does_not_exist(self.pycache_dir)
+
+    def test_that_distdir_is_removed(self):
+        self.assert_path_does_not_exist(self.dist_dir)
+
+    def test_that_eggdir_is_removed(self):
+        self.assert_path_does_not_exist(self.egg_dir)
+
+    def test_that_envdir_is_removed(self):
+        self.assert_path_does_not_exist(self.env_dir)