Updated Cython dev shim

CRNS/__init__.py

@@ -8,44 +8,47 @@ import threading
 from importlib.util import find_spec
 if find_spec("pyximport") and find_spec("cython"):
     import json
+    import shutil
+    import subprocess
     from importlib.abc import MetaPathFinder
     from importlib.machinery import ExtensionFileLoader, SourceFileLoader
     from importlib.machinery import EXTENSION_SUFFIXES
     from importlib.util import spec_from_file_location
     from distutils.extension import Extension
-
     from pyximport import pyxbuild
 
-    class _RNSImportFinder(MetaPathFinder):
+    COMPILE_DIRECTIVES = { "annotation_typing": False }
-        """A cythonizing import hook for the RNS package tree.
 
-        pyximport's own ``.py`` hook is unsuitable for this: it only
+    _WORKER_SCRIPT = r"""
-        resolves top-level module names (so every sub-module import
+import sys, os, json
-        misses the file), and it disables itself after compiling a
+from distutils.extension import Extension
-        single module.
+from pyximport import pyxbuild
+def main():
+    tasks_path = sys.argv[1]
+    with open(tasks_path, "r") as fh: tasks = json.load(fh)
+    results = []
+    for task in tasks:
+        try:
+            os.makedirs(task["scratch"], exist_ok=True)
+            extension = Extension(name=task["basename"], sources=[task["source"]])
+            extension.cython_directives = task["directives"]
+            so_path = pyxbuild.pyx_to_dll(task["source"], ext=extension, build_in_temp=True, pyxbuild_dir=task["scratch"])
+            os.makedirs(os.path.dirname(task["target"]), exist_ok=True)
+            os.replace(so_path, task["target"])
+            results.append({"fullname": task["fullname"], "ok": True, "error": None})
+        except Exception as e:
+            results.append({"fullname": task["fullname"], "ok": False,
+                            "error": "{}: {}".format(type(e).__name__, e)})
+    with open(tasks_path + ".results", "w") as fh: json.dump(results, fh)
+if __name__ == "__main__": main()
+"""
 
-        This finder compiles RNS source modules into a mirror tree
+    def _compile_jobs():
-        under ``~/.pyxbld/rnslib/`` that exactly mirrors the source
+        jobs = os.environ.get("CRNS_JOBS")
-        package layout, so that:
+        if jobs: return max(1, int(jobs))
+        return max(1, min(os.cpu_count() or 1, 16))
 
-          * ``os.path.dirname(__file__)`` of a compiled package
+    class _ImportFinder(MetaPathFinder):
-            initialiser equals that package's own mirror directory,
-            keeping the ``__all__`` globs in package ``__init__``
-            modules correct when the modules are compiled.
-          * sub-module resolution through ``__path__`` finds the
-            compiled extensions in the mirror.
-          * each loaded module is a C extension (``.so``/``.pyd``).
-
-        The entire RNS source tree is compiled up front (on first use,
-        and incrementally on later uses), so the mirror is complete
-        before any RNS module is imported and imported modules can be
-        loaded straight from it.
-
-        Third-party and standard-library modules are left untouched,
-        and the ``RNS.vendor`` subtree is deliberately imported from
-        source. Modules that cannot be compiled fall back to the
-        interpreted source module with a warning.
-        """
 
         def __init__(self):
             self.build_dir = os.path.join(os.path.expanduser("~"), ".pyxbld")
@@ -54,26 +57,40 @@ if find_spec("pyximport") and find_spec("cython"):
             self.abi_suffix = EXTENSION_SUFFIXES[0]
             self.unbuildable = {}
             self.unbuildable_path = os.path.join(self.build_dir, "unbuildable.json")
+            self._verify_mirror_directives()
             try:
-                with open(self.unbuildable_path, "r") as fh:
+                with open(self.unbuildable_path, "r") as fh: self.unbuildable = json.load(fh)
-                    self.unbuildable = json.load(fh)
+            except Exception: self.unbuildable = {}
-            except Exception:
+
-                self.unbuildable = {}
+        def _verify_mirror_directives(self):
+            # If the compiler directives have changed since the mirror was
+            # built, the entire mirror is invalidated and rebuilt.
+            stamp_path = os.path.join(self.build_dir, "rnslib.stamp")
+            directives = json.dumps(COMPILE_DIRECTIVES, sort_keys=True)
+            stamp = directives + "\n" + self.abi_suffix
+            previous = None
+            try:
+                with open(stamp_path, "r") as fh: previous = fh.read()
+            except Exception: pass
+            if previous != stamp:
+                shutil.rmtree(self.mirror_root, ignore_errors=True)
+                os.makedirs(self.mirror_root, exist_ok=True)
+                try:
+                    os.makedirs(self.build_dir, exist_ok=True)
+                    with open(stamp_path, "w") as fh: fh.write(stamp)
+                except Exception: pass
 
         def _save_unbuildable(self):
             try:
                 os.makedirs(os.path.dirname(self.unbuildable_path), exist_ok=True)
-                with open(self.unbuildable_path, "w") as fh:
+                with open(self.unbuildable_path, "w") as fh: json.dump(self.unbuildable, fh, indent=1)
-                    json.dump(self.unbuildable, fh, indent=1)
+            except Exception: pass
-            except Exception:
-                pass
 
         def _mark_unbuildable(self, fullname, source):
             try:
                 self.unbuildable[fullname] = int(os.path.getmtime(source))
                 self._save_unbuildable()
-            except Exception:
+            except Exception: pass
-                pass
 
         def _unmark_buildable(self, fullname):
             if fullname in self.unbuildable:
@@ -81,76 +98,65 @@ if find_spec("pyximport") and find_spec("cython"):
                 self._save_unbuildable()
 
         def _is_unbuildable(self, fullname, source):
-            if fullname not in self.unbuildable:
+            if fullname not in self.unbuildable: return False
-                return False
             try:
-                if int(os.path.getmtime(source)) == self.unbuildable[fullname]:
+                if int(os.path.getmtime(source)) == self.unbuildable[fullname]: return True
-                    return True
+            except Exception: pass
-            except Exception:
-                pass
             self.unbuildable.pop(fullname, None)
             return False
 
         def _find_source_root(self):
             here = os.path.dirname(os.path.abspath(__file__))
             candidate = os.path.abspath(os.path.join(here, "..", "RNS", "__init__.py"))
-            if os.path.exists(candidate):
+            if os.path.exists(candidate): return os.path.dirname(os.path.dirname(candidate))
-                return os.path.dirname(os.path.dirname(candidate))
             for entry in [os.getcwd()] + list(sys.path):
-                if not entry:
+                if not entry: continue
-                    continue
                 candidate = os.path.join(entry, "RNS", "__init__.py")
-                if os.path.exists(candidate):
+                if os.path.exists(candidate): return os.path.dirname(candidate)
-                    return os.path.dirname(candidate)
             raise RuntimeError("CRNS could not locate the RNS source tree")
 
         def _all_source_modules(self):
             rns_root = os.path.join(self.source_root, "RNS")
             for root, dirs, files in os.walk(rns_root):
-                dirs[:] = [d for d in dirs if d not in ("__pycache__", "vendor")]
+                dirs[:] = [d for d in dirs if d not in ("__pycache__", "i2plib")]
                 rel = os.path.relpath(root, self.source_root).replace(os.sep, ".")
                 for fn in sorted(files):
-                    if not fn.endswith(".py"):
+                    if not fn.endswith(".py"): continue
-                        continue
+                    if fn == "__init__.py": fullname = rel
-                    if fn == "__init__.py":
+                    else: fullname = rel + "." + fn[:-3]
-                        fullname = rel
-                    else:
-                        fullname = rel + "." + fn[:-3]
                     yield fullname
 
         def _mirror_target(self, fullname, is_package):
             rel = fullname.split(".")
             basename = rel[-1]
-            if is_package:
+            if is_package: return os.path.join(self.mirror_root, *rel, "__init__" + self.abi_suffix)
-                return os.path.join(self.mirror_root, *rel, "__init__" + self.abi_suffix)
             return os.path.join(self.mirror_root, *rel[:-1], basename + self.abi_suffix)
 
         def _source(self, fullname):
             rel = fullname.split(".")
             package_source = os.path.join(self.source_root, *rel, "__init__.py")
-            if os.path.exists(package_source):
+            if os.path.exists(package_source): return package_source, True
-                return package_source, True
             module_source = os.path.join(self.source_root, *rel) + ".py"
-            if os.path.exists(module_source):
+            if os.path.exists(module_source): return module_source, False
-                return module_source, False
             return None, None
 
         def _compile(self, fullname, source):
             basename = fullname.rsplit(".", 1)[-1]
             extension = Extension(name=basename, sources=[source])
-            so_path = pyxbuild.pyx_to_dll(
+            extension.cython_directives = COMPILE_DIRECTIVES
-                source,
+            so_path = pyxbuild.pyx_to_dll(source, ext=extension, build_in_temp=True, pyxbuild_dir=self.build_dir)
-                ext=extension,
-                build_in_temp=True,
-                pyxbuild_dir=self.build_dir,
-            )
             abi_suffix = os.path.basename(so_path)[len(basename):]
             return so_path, abi_suffix
 
+        def _is_current(self, fullname, source, is_package):
+            target = self._mirror_target(fullname, is_package)
+            if not os.path.exists(target): return False
+            try: return os.path.getmtime(target) >= os.path.getmtime(source)
+            except Exception: return False
+
         def _ensure_compiled(self, fullname, source, is_package):
             target = self._mirror_target(fullname, is_package)
-            if os.path.exists(target):
+            if os.path.exists(target) and os.path.getmtime(target) >= os.path.getmtime(source): return target
-                return target
             so_path, abi_suffix = self._compile(fullname, source)
             os.makedirs(os.path.dirname(target), exist_ok=True)
             os.replace(so_path, target)
@@ -159,74 +165,114 @@ if find_spec("pyximport") and find_spec("cython"):
         def precompile(self):
             if not os.path.isdir(self.mirror_root):
                 os.makedirs(self.mirror_root, exist_ok=True)
+
+            pending = []
             for fullname in self._all_source_modules():
                 source, is_package = self._source(fullname)
                 if source is None:
                     continue
-                if os.path.exists(self._mirror_target(fullname, is_package)):
+                if self._is_current(fullname, source, is_package):
                     self._unmark_buildable(fullname)
                     continue
                 if self._is_unbuildable(fullname, source):
                     continue
-                try:
+                
-                    self._ensure_compiled(fullname, source, is_package)
+                pending.append((fullname, source, is_package))
-                    self._unmark_buildable(fullname)
+
-                except Exception as e:
+            if not pending: return
-                    self._mark_unbuildable(fullname, source)
+
-                    print(f"CRNS: could not compile {fullname}, "
+            jobs = min(_compile_jobs(), len(pending))
-                          f"using interpreted module: {e}", file=sys.stderr)
+            if jobs > 1:
+                # Round-robin assignment spreads the large modules evenly
+                # across workers.
+                chunks = [[] for _ in range(jobs)]
+                for i, item in enumerate(pending): chunks[i % jobs].append(item)
+
+                workers_dir = os.path.join(self.build_dir, "workers")
+                os.makedirs(workers_dir, exist_ok=True)
+                procs = []
+                for w in range(jobs):
+                    tasks = []
+                    for fullname, source, is_package in chunks[w]:
+                        tasks.append({ "fullname": fullname,
+                                       "basename": fullname.rsplit(".", 1)[-1],
+                                       "source": source,
+                                       "target": self._mirror_target(fullname, is_package),
+                                       "scratch": os.path.join(workers_dir, "w{}".format(w)),
+                                       "directives": COMPILE_DIRECTIVES })
+                    task_path = os.path.join(workers_dir, "tasks{}.json".format(w))
+                    with open(task_path, "w") as fh: json.dump(tasks, fh)
+                    procs.append((w, task_path, subprocess.Popen([sys.executable, "-c", _WORKER_SCRIPT, task_path],
+                                                                 stdout=subprocess.PIPE, stderr=subprocess.PIPE)))
+
+                for w, task_path, proc in procs:
+                    stdout, stderr = proc.communicate()
+                    if proc.returncode != 0:
+                        for fullname, _, _ in chunks[w]:
+                            source, _ = self._source(fullname)
+                            self._mark_unbuildable(fullname, source)
+                            print(f"CRNS: compile worker failed for {fullname}, using interpreted module", file=sys.stderr)
+                        if stderr:
+                            tail = stderr.decode(errors="replace").strip().splitlines()[-5:]
+                            print("CRNS: worker output:\n" + "\n".join(tail), file=sys.stderr)
+                        continue
+                    
+                    results = {}
+                    try:
+                        with open(task_path + ".results", "r") as fh:
+                            for entry in json.load(fh): results[entry["fullname"]] = entry
+                    except Exception: results = {}
+                    
+                    for fullname, _, _ in chunks[w]:
+                        entry = results.get(fullname)
+                        if entry is not None and entry.get("ok"): self._unmark_buildable(fullname)
+                        else:
+                            error = "" if entry is None else entry.get("error", "")
+                            source, _ = self._source(fullname)
+                            self._mark_unbuildable(fullname, source)
+                            print(f"Could not compile {fullname}, using interpreted module: {error}", file=sys.stderr)
+            else:
+                for fullname, source, is_package in pending:
+                    try:
+                        self._ensure_compiled(fullname, source, is_package)
+                        self._unmark_buildable(fullname)
+                    except Exception as e:
+                        self._mark_unbuildable(fullname, source)
+                        print(f"Could not compile {fullname}, using interpreted module: {e}", file=sys.stderr)
 
         def _spec(self, fullname, source, is_package, compile_attempted):
-            if fullname == "RNS.vendor" or fullname.startswith("RNS.vendor."):
+            if fullname == "RNS.vendor.i2plib" or fullname.startswith("RNS.vendor.i2plib."):
                 rel = fullname.split(".")
                 locations = [os.path.join(self.source_root, *rel)] if is_package else None
-                return spec_from_file_location(
+                return spec_from_file_location(fullname, source, loader=SourceFileLoader(fullname, source), submodule_search_locations=locations)
-                    fullname, source,
-                    loader=SourceFileLoader(fullname, source),
-                    submodule_search_locations=locations,
-                )
 
             if compile_attempted:
                 try:
                     target = self._ensure_compiled(fullname, source, is_package)
                     self._unmark_buildable(fullname)
-                    if is_package:
+                    if is_package: locations = [os.path.dirname(target)]
-                        locations = [os.path.dirname(target)]
+                    else:          locations = None
-                    else:
+                    return spec_from_file_location(fullname, target, loader=ExtensionFileLoader(fullname, target),
-                        locations = None
+                                                   submodule_search_locations=locations)
-                    return spec_from_file_location(
-                        fullname, target,
-                        loader=ExtensionFileLoader(fullname, target),
-                        submodule_search_locations=locations,
-                    )
                 except Exception as e:
                     self._mark_unbuildable(fullname, source)
-                    print(f"CRNS: could not compile {fullname}, "
+                    print(f"Could not compile {fullname}, using interpreted module: {e}", file=sys.stderr)
-                          f"using interpreted module: {e}", file=sys.stderr)
 
             rel = fullname.split(".")
             locations = [os.path.join(self.source_root, *rel)] if is_package else None
-            return spec_from_file_location(
+            return spec_from_file_location(fullname, source,loader=SourceFileLoader(fullname, source),submodule_search_locations=locations)
-                fullname, source,
-                loader=SourceFileLoader(fullname, source),
-                submodule_search_locations=locations,
-            )
 
         def find_spec(self, fullname, path, target=None):
-            if not isinstance(fullname, str):
+            if not isinstance(fullname, str):                          return None
-                return None
+            if not (fullname == "RNS" or fullname.startswith("RNS.")): return None
-            if not (fullname == "RNS" or fullname.startswith("RNS.")):
-                return None
 
             source, is_package = self._source(fullname)
-            if source is None:
+            if source is None: return None
-                return None
 
-            return self._spec(fullname, source, is_package,
+            return self._spec(fullname, source, is_package, compile_attempted=not self._is_unbuildable(fullname, source))
-                              compile_attempted=not self._is_unbuildable(fullname, source))
 
-    _RNSFinder = _RNSImportFinder()
+    _importfinder = _ImportFinder()
-    sys.meta_path.insert(0, _RNSFinder)
+    sys.meta_path.insert(0, _importfinder)
 
 def notice_job():
     global noticed
@@ -234,14 +280,13 @@ def notice_job():
     while compiling:
         if time.time() > started+notice_delay and compiling:
             noticed = True
-            print("Compiling RNS object code... ", end="")
+            print("Compiling RNS object code...")
             sys.stdout.flush()
             break
         time.sleep(0.1)
 
-
 compiling = True
 threading.Thread(target=notice_job, daemon=True).start()
-_RNSFinder.precompile()
+_importfinder.precompile()
 import RNS; compiling = False
-if noticed: print("Done."); sys.stdout.flush()
+if noticed: print("Compilation done"); sys.stdout.flush()