Modify the generated CLI script by monkey patching

This change monkey patches setuptools' code generation functionality, so
that pkg_resources won't be imported in the generated file.

Fixes #172

1 files changed, 37 insertions, 0 deletions

diff --git a/setup.py b/setup.py
index c98d65136..3e9f88545 100755
--- a/setup.py
+++ b/setup.py
@@ -28,6 +28,7 @@ if sys.version_info[0] != 3 or sys.version_info[1] < 4:
 
 try:
     from setuptools import setup, find_packages
+    from setuptools.command.easy_install import ScriptWriter
 except ImportError:
     print("BuildStream requires setuptools in order to build. Install it using"
           " your package manager (usually python3-setuptools) or via pip (pip3"
@@ -134,6 +135,42 @@ if not os.environ.get('BST_ARTIFACTS_ONLY', ''):
 
 
 #####################################################
+#    Monkey-patching setuptools for performance     #
+#####################################################
+#
+# The template of easy_install.ScriptWriter is inefficient in our case as it
+# imports pkg_resources. Patching the template only doesn't work because of the
+# old string formatting used (%). This forces us to overwrite the class function
+# as well.
+#
+# The patch was inspired from https://github.com/ninjaaron/fast-entry_points
+# which we believe was also inspired from the code from `setuptools` project.
+TEMPLATE = '''\
+# -*- coding: utf-8 -*-
+import sys
+
+from {0} import {1}
+
+if __name__ == '__main__':
+    sys.exit({2}())'''
+
+
+@classmethod
+def get_args(cls, dist, header=None):
+    if header is None:
+        header = cls.get_header()
+    for name, ep in dist.get_entry_map('console_scripts').items():
+        cls._ensure_safe_name(name)
+        script_text = TEMPLATE.format(ep.module_name, ep.attrs[0], '.'.join(ep.attrs))
+        args = cls._get_script_args('console', name, header, script_text)
+        for res in args:
+            yield res
+
+
+ScriptWriter.get_args = get_args
+
+
+#####################################################
 #             Main setup() Invocation               #
 #####################################################
 setup(name='BuildStream',