implement a custom 'python setup.py clean' command which gets rid of junk files

1 files changed, 22 insertions, 0 deletions

diff --git a/setup.py b/setup.py
index 17e060b5..70f1d1cc 100644
--- a/setup.py
+++ b/setup.py
@@ -6,12 +6,30 @@
 
 import sys
 import os
+import shutil
+import fnmatch
 try:
     from setuptools import setup, Extension
 except ImportError:
     from distutils.core import setup, Extension
 
 
+def clean():
+    """'python setup.py clean' custom command."""
+    def rglob(pattern):
+        return [os.path.join(dirpath, f)
+            for dirpath, dirnames, files in os.walk('.')
+            for f in fnmatch.filter(files, pattern)]
+
+    for dirpath, dirnames, files in os.walk('.'):
+        if dirpath.endswith('__pycache__'):
+            sys.stdout.write('removing %s\n' % dirpath)
+            shutil.rmtree(dirpath)
+    for pattern in ['*.py[co]', '*~', '*.orig', '*.rej', '*.swp']:
+        for x in rglob(pattern):
+           sys.stdout.write('removing %s\n' % x)
+           os.remove(x)
+
 def get_version():
     INIT = os.path.abspath(os.path.join(os.path.dirname(__file__),
                            'psutil', '__init__.py'))
@@ -102,6 +120,10 @@ else:
 
 
 def main():
+    # "python setup.py clean" custom command
+    if len(sys.argv) > 1 and sys.argv[1] == 'clean':
+        return clean()
+
     setup_args = dict(
         name='psutil',
         version=VERSION,