Merge pull request #2 from dave-shawley/add-egg-support

Add egg support

Author: dave-shawley

HISTORY

@@ -1,6 +1,11 @@
 Changelog
 =========
 
-* 0.0.1
+* 0.0.2 (17-Nov-2014)
+
+  - Support removal of .egg-info and .egg directories.
+  - Add support for *--dry-run*
+
+* 0.0.1 (16-Nov-2014)
 
   - Support removal of distribution directories.

README.rst

@@ -67,6 +67,9 @@ few new command line parameters.
 ``setup.py clean --dist``
    Removes directories that the various *dist* commands produce.
 
+``setup.py clean --egg``
+   Removes *.egg* and *.egg-info* directories.
+
 Where can I get this extension from?
 ------------------------------------
 +---------------+-----------------------------------------------------+
@@ -85,11 +88,11 @@ Where can I get this extension from?
    #extending-and-reusing-setuptools
 .. _setuptools: https://pythonhosted.org/setuptools/
 
-.. |Version| image:: https://badge.fury.io/py/setupext-janitor.svg
+.. |Version| image:: https://badge.fury.io/py/setupext-janitor.svg?
    :target: https://badge.fury.io/
 .. |Downloads| image:: https://pypip.in/d/setupext-janitor/badge.svg?
    :target: https://pypi.python.org/pypi/setupext-janitor
 .. |Status| image:: https://travis-ci.org/dave-shawley/setupext-janitor.svg
    :target: https://travis-ci.org/dave-shawley/setupext-janitor
-.. |License| image:: https://pypip.in/license/dave-shawley/badge.svg?
-   :target: https://setupext-dave-shawley.readthedocs.org/
+.. |License| image:: https://pypip.in/license/setupext-janitor/badge.svg?
+   :target: https://setupext-janitor.readthedocs.org/

setup.py

@@ -44,4 +44,7 @@
             'clean = setupext.janitor:CleanCommand',
         ],
     },
+    cmdclass={
+        'clean': janitor.CleanCommand,
+    },
 )

setupext/janitor/__init__.py

@@ -1,9 +1,9 @@
-from distutils import log
+from distutils import dir_util, errors
 from distutils.command.clean import clean as _CleanCommand
-import shutil
+import os.path
 
 
-version_info = (0, 0, 1)
+version_info = (0, 0, 2)
 __version__ = '.'.join(str(v) for v in version_info)
 
 
@@ -36,30 +36,49 @@ class CleanCommand(_CleanCommand):
 
     # See _set_options for `user_options`
 
-    def __init__(self, *args, **kwargs):
-        _CleanCommand.__init__(self, *args, **kwargs)
-        self.dist = None
-
     def initialize_options(self):
         _CleanCommand.initialize_options(self)
         self.dist = False
+        self.eggs = False
+        self.egg_base = None
+
+    def finalize_options(self):
+        _CleanCommand.finalize_options(self)
+        try:
+            self.set_undefined_options(
+                'egg_info', ('egg_base', 'egg_base'))
+        except errors.DistutilsError:
+            pass
+
+        if self.egg_base is None:
+            self.egg_base = os.curdir
 
     def run(self):
         _CleanCommand.run(self)
-        if not self.dist:
-            return
 
-        dist_dirs = set()
-        for cmd_name in self.distribution.commands:
-            if 'dist' in cmd_name:
-                command = self.distribution.get_command_obj(cmd_name)
-                command.ensure_finalized()
-                if getattr(command, 'dist_dir', None):
-                    dist_dirs.add(command.dist_dir)
-
-        for dir_name in dist_dirs:
-            self.announce('removing {0}'.format(dir_name), level=log.DEBUG)
-            shutil.rmtree(dir_name, ignore_errors=True)
+        dir_names = set()
+        if self.dist:
+            for cmd_name, _ in self.distribution.get_command_list():
+                if 'dist' in cmd_name:
+                    command = self.distribution.get_command_obj(cmd_name)
+                    command.ensure_finalized()
+                    if getattr(command, 'dist_dir', None):
+                        dir_names.add(command.dist_dir)
+
+        if self.eggs:
+            for name in os.listdir(self.egg_base):
+                if name.endswith('.egg-info'):
+                    dir_names.add(os.path.join(self.egg_base, name))
+            for name in os.listdir(os.curdir):
+                if name.endswith('.egg'):
+                    dir_names.add(name)
+
+        for dir_name in dir_names:
+            if os.path.exists(dir_name):
+                dir_util.remove_tree(dir_name, dry_run=self.dry_run)
+            else:
+                self.announce(
+                    'skipping {0} since it does not exist'.format(dir_name))
 
 
 def _set_options():
@@ -81,8 +100,13 @@ def _set_options():
     CleanCommand.user_options = _CleanCommand.user_options[:]
     CleanCommand.user_options.extend([
         ('dist', 'd', 'remove distribution directory'),
+        ('eggs', None, 'remove egg and egg-info directories'),
+
+        ('egg-base=', 'e',
+         'directory containing .egg-info directories '
+         '(default: top of the source tree)'),
     ])
     CleanCommand.boolean_options = _CleanCommand.boolean_options[:]
-    CleanCommand.boolean_options.append('dist')
+    CleanCommand.boolean_options.extend(['dist', 'eggs'])
 
 _set_options()

tests.py

@@ -5,6 +5,7 @@
 import shutil
 import sys
 import tempfile
+import uuid
 
 if sys.version_info >= (2, 7):
     import unittest
@@ -71,26 +72,31 @@ def test_that_janitor_defines_dist_command(self):
             janitor.CleanCommand.user_options)
 
 
-class DirectoryCleanupTests(unittest.TestCase):
-    temp_dir = tempfile.mkdtemp()
+class DirectoryCleanupMixin(object):
 
     @classmethod
     def setUpClass(cls):
-        super(DirectoryCleanupTests, cls).setUpClass()
+        super(DirectoryCleanupMixin, cls).setUpClass()
+        cls.temp_dir = tempfile.mkdtemp()
         atexit.register(shutil.rmtree, cls.temp_dir)
 
     @classmethod
     def create_directory(cls, dir_name):
         return tempfile.mkdtemp(dir=cls.temp_dir, prefix=dir_name)
 
-    def assert_path_does_not_exist(self, full_path):
+    def assert_path_does_not_exist(self, *trailing_segments):
+        full_path = os.path.join(*trailing_segments)
         if os.path.exists(full_path):
             raise AssertionError('{0} should not exist'.format(full_path))
 
-    def assert_path_exists(self, full_path):
+    def assert_path_exists(self, *trailing_segments):
+        full_path = os.path.join(*trailing_segments)
         if not os.path.exists(full_path):
             raise AssertionError('{0} should exist'.format(full_path))
 
+
+class DistDirectoryCleanupTests(DirectoryCleanupMixin, unittest.TestCase):
+
     def test_that_dist_directory_is_removed_for_sdist(self):
         dist_dir = self.create_directory('dist-dir')
         run_setup(
@@ -128,3 +134,50 @@ def test_that_directories_are_not_removed_without_parameter(self):
         )
         self.assert_path_exists(sdist_dir)
         self.assert_path_exists(bdist_dir)
+
+    def test_that_directories_are_not_removed_in_dry_run_mode(self):
+        sdist_dir = self.create_directory('sdist-dir')
+        run_setup(
+            'sdist', '--dist-dir={0}'.format(sdist_dir),
+            'clean', '--dist', '--dry-run'
+        )
+        self.assert_path_exists(sdist_dir)
+
+
+class EggDirectoryCleanupTests(DirectoryCleanupMixin, unittest.TestCase):
+
+    def test_that_egg_info_directories_are_removed(self):
+        egg_root = self.create_directory('egg-info-root')
+        os.mkdir(os.path.join(egg_root, 'bah.egg-info'))
+        os.mkdir(os.path.join(egg_root, 'foo.egg-info'))
+        run_setup('clean', '--egg-base={0}'.format(egg_root), '--eggs')
+        self.assert_path_exists(egg_root)
+        self.assert_path_does_not_exist(egg_root, 'bah.egg-info')
+        self.assert_path_does_not_exist(egg_root, 'foo.egg-info')
+
+    def test_that_egg_directories_are_removed(self):
+        dir_name = uuid.uuid4().hex + '.egg'
+        os.mkdir(dir_name)
+        try:
+            run_setup('clean', '--eggs')
+            self.assert_path_does_not_exist(dir_name)
+        except:
+            os.rmdir(dir_name)
+            raise
+
+    def test_that_directories_are_not_removed_in_dry_run_mode(self):
+        egg_root = self.create_directory('egg-info-root')
+        os.mkdir(os.path.join(egg_root, 'package.egg-info'))
+        installed_egg = uuid.uuid4().hex + '.egg'
+        os.mkdir(installed_egg)
+        try:
+            run_setup(
+                'clean', '--egg-base={0}'.format(egg_root), '--eggs',
+                '--dry-run',
+            )
+            self.assert_path_exists(installed_egg)
+            self.assert_path_exists(egg_root, 'package.egg-info')
+        except:
+            os.rmdir(installed_egg)
+            raise
+        os.rmdir(installed_egg)