From e181ba75e458f22a932963a414a7079edb85e410 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Piotr=20Ma=C5=9Blanka?= <piotr.maslanka@henrietta.com.pl>
Date: Wed, 9 Dec 2020 16:40:19 +0100
Subject: [PATCH] add `monkey_patch_parallel_compilation`

---
 CHANGELOG.md         |  1 +
 docs/distutils.rst   |  5 +++++
 docs/index.rst       |  1 +
 satella/__init__.py  |  2 +-
 satella/distutils.py | 43 +++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 51 insertions(+), 1 deletion(-)
 create mode 100644 docs/distutils.rst
 create mode 100644 satella/distutils.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5715081c..ecbc8b74 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,3 +3,4 @@
 * added `write_json_to_file`
 * added `read_json_from_file`
 * added `write_json_to_file_if_different`
+* added `satella.distutils`
diff --git a/docs/distutils.rst b/docs/distutils.rst
new file mode 100644
index 00000000..552c4b9d
--- /dev/null
+++ b/docs/distutils.rst
@@ -0,0 +1,5 @@
+Distutils extensions
+====================
+
+.. autofunction:: satella.distutils.monkey_patch_parallel_compilation
+
diff --git a/docs/index.rst b/docs/index.rst
index 7a13933c..a7eaa31a 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -38,6 +38,7 @@ Visit the project's page at GitHub_!
            processes
            cassandra
            opentracing
+           distutils
 
 
 Indices and tables
diff --git a/satella/__init__.py b/satella/__init__.py
index 702d0b12..16d5df04 100644
--- a/satella/__init__.py
+++ b/satella/__init__.py
@@ -1 +1 @@
-__version__ = '2.14.23_a2'
+__version__ = '2.14.23_a3'
diff --git a/satella/distutils.py b/satella/distutils.py
new file mode 100644
index 00000000..0a048b96
--- /dev/null
+++ b/satella/distutils.py
@@ -0,0 +1,43 @@
+import typing as tp
+import multiprocessing
+
+__all__ = ['monkey_patch_parallel_compilation']
+
+
+def monkey_patch_parallel_compilation(cores: tp.Optional[int] = None):
+    """
+    This monkey-patches distutils to provide parallel compilation, even if you have
+    a single extension built from multiple .c files.
+
+    Invoke in your setup.py file
+
+    :param cores: amount of cores. Leave at default (None) for autodetection.
+    """
+    if cores is None:
+        cores = multiprocessing.cpu_count()
+
+    # monkey-patch for parallel compilation
+    def parallelCCompile(self, sources, output_dir=None, macros=None, include_dirs=None, debug=0,
+                         extra_preargs=None, extra_postargs=None, depends=None):
+        # those lines are copied from distutils.ccompiler.CCompiler directly
+        macros, objects, extra_postargs, pp_opts, build = self._setup_compile(output_dir, macros,
+                                                                              include_dirs, sources,
+                                                                              depends,
+                                                                              extra_postargs)
+        cc_args = self._get_cc_args(pp_opts, debug, extra_preargs)
+        # parallel code
+        N = 2  # number of parallel compilations
+        import multiprocessing.pool
+        def _single_compile(obj):
+            try:
+                src, ext = build[obj]
+            except KeyError:
+                return
+            self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
+
+        # convert to list, imap is evaluated on-demand
+        list(multiprocessing.pool.ThreadPool(cores).imap(_single_compile, objects))
+        return objects
+
+    import distutils.ccompiler
+    distutils.ccompiler.CCompiler.compile = parallelCCompile
-- 
GitLab