optimize_upip.py: Script to optimize archives for low-heap upip usage.

optimize_upip.py

@@ -0,0 +1,106 @@
+#
+# This script optimizes a Python source distribution tarball as produced by
+# "python3 setup.py sdist" command for MicroPython's native package manager,
+# upip. Optimization includes:
+#  * Removing metadata files not used by upip (this includes setup.py)
+#  * Recompressing gzip archive with 4K dictionary size so it can be
+#    installed even on low-heap targets.
+#
+import sys
+import os
+import zlib
+from subprocess import Popen, PIPE
+import glob
+import tarfile
+import re
+import io
+
+
+def gzip_4k(inf, fname):
+    comp = zlib.compressobj(level=9, wbits=16 + 12)
+    with open(fname + ".out", "wb") as outf:
+        while 1:
+            data = inf.read(1024)
+            if not data:
+                break
+            outf.write(comp.compress(data))
+        outf.write(comp.flush())
+    os.rename(fname, fname + ".orig")
+    os.rename(fname + ".out", fname)
+
+
+def recompress(fname):
+    with Popen(["gzip", "-d", "-c", fname], stdout=PIPE).stdout as inf:
+        gzip_4k(inf, fname)
+
+def find_latest(dir):
+    res = []
+    for fname in glob.glob(dir + "/*.gz"):
+        st = os.stat(fname)
+        res.append((st.st_mtime, fname))
+    res.sort()
+    latest = res[-1][1]
+    return latest
+
+
+def recompress_latest(dir):
+    latest = find_latest(dir)
+    print(latest)
+    recompress(latest)
+
+
+EXCLUDE = [r".+/setup.py"]
+INCLUDE = [r".+\.py", r".+\.egg-info/(PKG-INFO|requires\.txt)"]
+
+
+outbuf = io.BytesIO()
+
+def filter_tar(name):
+    fin = tarfile.open(name, "r:gz")
+    fout = tarfile.open(fileobj=outbuf, mode="w")
+    for info in fin:
+#        print(info)
+        include = None
+        for p in EXCLUDE:
+            if re.match(p, info.name):
+                include = False
+                break
+        if include is None:
+            for p in INCLUDE:
+                if re.match(p, info.name):
+                    include = True
+                    print("Including:", info.name)
+        if not include:
+            continue
+        farch = fin.extractfile(info)
+        fout.addfile(info, farch)
+    fout.close()
+    fin.close()
+
+
+
+from setuptools import Command
+
+class OptimizeUpip(Command):
+
+    user_options = []
+
+    def run(self):
+        latest = find_latest("dist")
+        filter_tar(latest)
+        outbuf.seek(0)
+        gzip_4k(outbuf, latest)
+
+    def initialize_options(self):
+        pass
+
+    def finalize_options(self):
+        pass
+
+
+# For testing only
+if __name__ == "__main__":
+#    recompress_latest(sys.argv[1])
+    filter_tar(sys.argv[1])
+    outbuf.seek(0)
+    gzip_4k(outbuf, sys.argv[1])