Consider all installed files for cleaning up

src/sage_setup/clean.py

@@ -57,7 +57,7 @@ def _remove(file_set, module_base, to_remove):
         file_set.difference_update([filename])
 
 
-def _find_stale_files(site_packages, python_packages, python_modules, ext_modules):
+def _find_stale_files(site_packages, python_packages, python_modules, ext_modules, data_files):
     """
     Find stale files
 
@@ -70,16 +70,18 @@ def _find_stale_files(site_packages, python_packages, python_modules, ext_module
     course. We check that when the doctest is being run, that is,
     after installation, there are no stale files::
 
-        sage: from sage.env import SAGE_SRC, SAGE_LIB
-        sage: from sage_setup.find import find_python_sources
+        sage: from sage.env import SAGE_SRC, SAGE_LIB, SAGE_CYTHONIZED
+        sage: from sage_setup.find import find_python_sources, find_extra_files
         sage: python_packages, python_modules = find_python_sources(
         ....:     SAGE_SRC, ['sage', 'sage_setup'])
+        sage: extra_files = find_extra_files(python_packages, SAGE_SRC,
+        ....:     SAGE_CYTHONIZED, SAGE_LIB, ["ntlwrap.cpp"])
         sage: from sage_setup.clean import _find_stale_files
 
     TODO: move ``module_list.py`` into ``sage_setup`` and also check
     extension modules::
 
-        sage: stale_iter = _find_stale_files(SAGE_LIB, python_packages, python_modules, [])
+        sage: stale_iter = _find_stale_files(SAGE_LIB, python_packages, python_modules, [], extra_files)
         sage: from sage.misc.sageinspect import loadable_module_extension
         sage: skip_extensions = (loadable_module_extension(),)
         sage: for f in stale_iter:
@@ -90,7 +92,6 @@ def _find_stale_files(site_packages, python_packages, python_modules, ext_module
     PYMOD_EXTS = (os.path.extsep + 'py', os.path.extsep + 'pyc')
     CEXTMOD_EXTS = (loadable_module_extension(),)
     INIT_FILES = tuple('__init__' + x for x in PYMOD_EXTS)
-    ALLOW_EXTS = ('.c', '.cpp', '.h', '.pxi', '.pxd')
 
     module_files = installed_files_by_module(site_packages, ['sage', 'sage_setup'])
 
@@ -115,13 +116,22 @@ def _find_stale_files(site_packages, python_packages, python_modules, ext_module
             continue
         _remove(files, mod, CEXTMOD_EXTS)
 
+    # Convert data_files to a set
+    installed_files = set()
+    for dir, files in data_files:
+        dir = os.path.relpath(dir, site_packages)
+        if dir.startswith("."):  # dir is not inside site_packages
+            continue
+        for f in files:
+            installed_files.add(os.path.join(dir, os.path.basename(f)))
+
     for files in module_files.values():
         for f in files:
-            if not f.endswith(ALLOW_EXTS):
+            if f not in installed_files:
                 yield f
 
 
-def clean_install_dir(site_packages, python_packages, python_modules, ext_modules):
+def clean_install_dir(site_packages, python_packages, python_modules, ext_modules, data_files):
     """
     Delete all modules that are **not** being installed
 
@@ -144,9 +154,12 @@ def clean_install_dir(site_packages, python_packages, python_modules, ext_module
 
     - ``ext_modules`` -- list of distutils ``Extension`` classes. The
       output of ``cythonize``.
+
+    - ``data_files`` -- a list of (installation directory, files) pairs,
+      like the ``data_files`` argument to distutils' ``setup()``.
     """
     stale_file_iter = _find_stale_files(
-        site_packages, python_packages, python_modules, ext_modules)
+        site_packages, python_packages, python_modules, ext_modules, data_files)
     for f in stale_file_iter:
         f = os.path.join(site_packages, f)
         print('Cleaning up stale file: {0}'.format(f))

src/sage_setup/find.py

@@ -74,6 +74,59 @@ def find_python_sources(src_dir, modules=('sage',)):
     return python_packages, python_modules
 
 
+def find_extra_files(packages, src_dir, cythonized_dir, site_packages, special_filenames=[]):
+    """
+    Find all extra files which should be installed.
+
+    These are:
+
+    1. From ``src_dir``: all .pxd and .pxi files and files listed in
+       ``special_filenames``.
+    2. From ``cythonized_dir``: all .h files (these are both the .h files
+       from the Sage sources, as well as all Cython-generated .h files).
+
+    INPUT:
+
+    - ``packages`` -- a list of Python packages to be considered
+
+    - ``src_dir`` -- the directory where to look for source files
+
+    - ``cythonized_dir`` -- the directory where the Cython-generated
+      files are
+
+    - ``site_packages`` -- the directory where the files should be
+      installed
+
+    - ``special_filenames`` -- a list of filenames to be installed from
+      ``src_dir``
+
+    EXAMPLES::
+
+        sage: from sage_setup.find import find_extra_files
+        sage: from sage.env import SAGE_SRC, SAGE_CYTHONIZED
+        sage: find_extra_files(["sage.ext.interrupt"], "src", SAGE_CYTHONIZED, ".")
+        [('./sage/ext/interrupt',
+          ['src/sage/ext/interrupt/interrupt.pxd', ...interrupt_api.h...])]
+    """
+    data_files = []
+
+    for package in packages:
+        dir = package.replace('.', os.path.sep)
+        sdir = os.path.join(src_dir, dir)
+        cydir = os.path.join(cythonized_dir, dir)
+
+        files = [os.path.join(sdir, f) for f in os.listdir(sdir)
+                if f.endswith((".pxd", ".pxi")) or f in special_filenames]
+        if os.path.isdir(cydir):  # Not every directory contains Cython files
+            files += [os.path.join(cydir, f) for f in os.listdir(cydir)
+                    if f.endswith(".h")]
+
+        if files:
+            data_files.append((os.path.join(site_packages, dir), files))
+
+    return data_files
+
+
 def installed_files_by_module(site_packages, modules=('sage',)):
     """
     Find all currently installed files

src/setup.py

@@ -606,39 +606,17 @@ def run_cythonize():
 
 
 #########################################################
-### Discovering Python Sources
+### Discovering Sources
 #########################################################
 
-print("Discovering Python source code....")
+print("Discovering Python/Cython source code....")
 t = time.time()
-from sage_setup.find import find_python_sources
+from sage_setup.find import find_python_sources, find_extra_files
 python_packages, python_modules = find_python_sources(
     SAGE_SRC, ['sage', 'sage_setup'])
-print("Discovered Python source, time: %.2f seconds." % (time.time() - t))
-
-
-#########################################################
-### Extra files to install
-#########################################################
-
-python_package_data = {}
-for package in python_packages:
-    if package == 'sage.libs.ntl':
-        python_package_data[package] = ['*.pxd', '*.pxi', '*.h', 'ntlwrap.cpp']
-    else:
-        python_package_data[package] = ['*.pxd', '*.pxi', '*.h']
-
-# List of files generated by cython that needs to be installed.
-# The list gives the location of the files to be installed relative to SAGE_LIB.
-cython_generated_files =[
-    os.path.join('sage', 'ext', 'interrupt', 'interrupt_api.h'),
-    os.path.join('sage', 'ext', 'interrupt', 'interrupt.h')
-]
-
-python_data_files = [
-    (os.path.join(SAGE_LIB, os.path.dirname(file)), [os.path.join(SAGE_CYTHONIZED, file)])
-    for file in cython_generated_files
-]
+python_data_files = find_extra_files(python_packages,
+    ".", SAGE_CYTHONIZED, SAGE_LIB, ["ntlwrap.cpp"])
+print("Discovered Python/Cython sources, time: %.2f seconds." % (time.time() - t))
 
 
 #########################################################
@@ -651,7 +629,8 @@ def run_cythonize():
 output_dirs = SITE_PACKAGES + glob.glob(os.path.join(SAGE_SRC, 'build', 'lib*'))
 for output_dir in output_dirs:
     print('- cleaning {0}'.format(output_dir))
-    clean_install_dir(output_dir, python_packages, python_modules, ext_modules)
+    clean_install_dir(output_dir, python_packages, python_modules,
+            ext_modules, python_data_files)
 print('Finished cleaning, time: %.2f seconds.' % (time.time() - t))
 
 
@@ -667,7 +646,6 @@ def run_cythonize():
       author_email= 'http://groups.google.com/group/sage-support',
       url         = 'http://www.sagemath.org',
       packages    = python_packages,
-      package_data= python_package_data,
       data_files  = python_data_files,
       scripts = [],
       cmdclass = { 'build_ext': sage_build_ext },