removing openw class

Author: ye11owSub

create_shadertext.py

@@ -1,134 +1,113 @@
-import io as cStringIO
-
-import sys, os
+import sys
 import time
-import re
-import glob
 from collections import defaultdict
+from itertools import chain
 from os.path import dirname
-from subprocess import Popen, PIPE
+from pathlib import Path
+from subprocess import PIPE, Popen
 
-def create_all(generated_dir, pymoldir="."):
-    '''
-    Generate various stuff
-    '''
-    create_shadertext(
-            os.path.join(pymoldir, "data", "shaders"),
-            generated_dir,
-            os.path.join(generated_dir, "ShaderText.h"),
-            os.path.join(generated_dir, "ShaderText.cpp"))
-    create_buildinfo(generated_dir, pymoldir)
 
-class openw(object):
+def create_all(generated_dir: str, pymol_dir: str = "."):
     """
-    File-like object for writing files. File is actually only
-    written if the content changed.
+    Generate various stuff
     """
-    def __init__(self, filename):
-        if os.path.exists(filename):
-            self.out = cStringIO.StringIO()
-            self.filename = filename
-        else:
-            os.makedirs(os.path.dirname(filename), exist_ok=True)
-            self.out = open(filename, "w")
-            self.filename = None
-    def close(self):
-        if self.out.closed:
-            return
-        if self.filename:
-            with open(self.filename) as handle:
-                oldcontents = handle.read()
-            newcontents = self.out.getvalue()
-            if oldcontents != newcontents:
-                self.out = open(self.filename, "w")
-                self.out.write(newcontents)
-        self.out.close()
-    def __getattr__(self, name):
-        return getattr(self.out, name)
-    def __enter__(self):
-        return self
-    def __exit__(self, *a, **k):
-        self.close()
-    def __del__(self):
-        self.close()
-
-def create_shadertext(shaderdir, shaderdir2, outputheader, outputfile):
-
-    outputheader = openw(outputheader)
-    outputfile = openw(outputfile)
+    generated_dir_path = Path(generated_dir)
+    pymol_dir_path = Path(pymol_dir)
 
+    generated_dir_path.mkdir(parents=True, exist_ok=True)
+    pymol_dir_path.mkdir(parents=True, exist_ok=True)
+
+    create_shadertext(
+        shader_dir=generated_dir_path / "data" / "shaders",
+        shader_dir2=pymol_dir_path,
+        output_header=generated_dir_path / "ShaderText.h",
+        output_source=generated_dir_path / "ShaderText.cpp",
+    )
+    create_buildinfo(generated_dir, pymol_dir)
+
+
+def create_shadertext(
+    shader_dir: Path,
+    shader_dir2: Path,
+    output_header: Path,
+    output_source: Path,
+):
+    varname = "_shader_cache_raw"
     include_deps = defaultdict(set)
     ifdef_deps = defaultdict(set)
+    extension_regexp = {".gs", ".vs", ".fs", ".shared", ".tsc", ".tse"}
 
     # get all *.gs *.vs *.fs *.shared from the two input directories
-    shaderfiles = set()
-    for sdir in [shaderdir, shaderdir2]:
-        for ext in ['gs', 'vs', 'fs', 'shared', 'tsc', 'tse']:
-            shaderfiles.update(map(os.path.basename,
-                sorted(glob.glob(os.path.join(sdir, '*.' + ext)))))
-
-    varname = '_shader_cache_raw'
-    outputheader.write('extern const char * %s[];\n' % varname)
-    outputfile.write('const char * %s[] = {\n' % varname)
-
-    for filename in sorted(shaderfiles):
-        shaderfile = os.path.join(shaderdir, filename)
-        if not os.path.exists(shaderfile):
-            shaderfile = os.path.join(shaderdir2, filename)
-
-        with open(shaderfile, 'r') as handle:
-            contents = handle.read()
-
-        if True:
-            outputfile.write('"%s", ""\n' % (filename))
-
+    shader_files = set(
+        chain(
+            (path for path in shader_dir.glob("**/*") if path.suffix in extension_regexp),
+            (path for path in shader_dir2.glob("**/*") if path.suffix in extension_regexp)
+        )
+    )
+
+    with (
+        open(output_header, "w") as output_header_file,
+        open(output_source, "w") as output_source_file,
+    ):
+        output_header_file.write(f"extern const char * {varname}[];\n")
+        output_source_file.write(f"const char * {varname}[] = {{\n")
+
+        for shader_file in shader_files:
+            output_source_file.write(f'"{shader_file.name}", ""\n')
+
+            contents = shader_file.read_text()
             for line in contents.splitlines():
                 line = line.strip()
 
                 # skip blank lines and obvious comments
-                if not line or line.startswith('//') and not '*/' in line:
+                if not line or line.startswith("//") and not "*/" in line:
                     continue
 
                 # write line, quoted, escaped and with a line feed
-                outputfile.write("\"%s\\n\"\n" % line.replace('\\', '\\\\').replace('"', r'\"'))
+                escaped_line = line.replace("\\", "\\\\").replace('"', r"\"")
+                output_source_file.write(f'"{escaped_line}\\n"\n')
 
                 # include and ifdef dependencies
-                if line.startswith('#include'):
-                    include_deps[line.split()[1]].add(filename)
-                elif line.startswith('#ifdef') or line.startswith('#ifndef'):
-                    ifdef_deps[line.split()[1]].add(filename)
-
-            outputfile.write(',\n')
-
-    outputfile.write('0};\n')
-
-    # include and ifdef dependencies
-    for varname, deps in [
-            ('_include_deps', include_deps),
-            ('_ifdef_deps', ifdef_deps)]:
-        outputheader.write('extern const char * %s[];\n' % varname)
-        outputfile.write('const char * %s[] = {\n' % varname)
-        for name, itemdeps in deps.items():
-            outputfile.write('"%s", "%s", 0,\n' % (name, '", "'.join(sorted(itemdeps))))
-        outputfile.write('0};\n')
-
-    outputheader.close()
-    outputfile.close()
-
-def create_buildinfo(outputdir, pymoldir='.'):
-
-    try:
-        sha = Popen(['git', 'rev-parse', 'HEAD'], cwd=pymoldir,
-                stdout=PIPE).stdout.read().strip().decode()
-    except OSError:
-        sha = ''
-
-    with openw(os.path.join(outputdir, 'PyMOLBuildInfo.h')) as out:
-        print('''
-#define _PyMOL_BUILD_DATE %d
-#define _PYMOL_BUILD_GIT_SHA "%s"
-        ''' % (int(os.environ.get('SOURCE_DATE_EPOCH', time.time())), sha), file=out)
+                if line.startswith("#include"):
+                    include_deps[line.split()[1]].add(shader_file.name)
+                elif line.startswith("#ifdef") or line.startswith("#ifndef"):
+                    ifdef_deps[line.split()[1]].add(shader_file.name)
+
+            output_source_file.write(",\n")
+        output_source_file.write("0};\n")
+
+        # include and ifdef dependencies
+        for varname, deps in [
+            ("_include_deps", include_deps),
+            ("_ifdef_deps", ifdef_deps),
+        ]:
+            output_header_file.write(f"extern const char * {varname}[];\n")
+            output_source_file.write(f"const char * {varname}[] = {{\n")
+            for name, item_deps in deps.items():
+                item_deps = '", "'.join(sorted(item_deps))
+                output_source_file.write(f'"{name}", "{item_deps}", 0,\n')
+            output_source_file.write("0};\n")
+
+
+def create_buildinfo(output_dir_path: str, pymoldir: str = "."):
+    output_dir = Path(output_dir_path)
+    sha_raw = Popen(["git", "rev-parse", "HEAD"], cwd=pymoldir, stdout=PIPE).stdout
+    sha = sha_raw.read().strip().decode() if sha_raw is not None else ""
+
+    info_file = output_dir / "PyMOLBuildInfo.h"
+    info_file.write_text(
+        f"""
+    #define _PyMOL_BUILD_DATE {time.time()}
+    #define _PYMOL_BUILD_GIT_SHA "{sha}"
+    """
+    )
+
 
 if __name__ == "__main__":
-    create_shadertext(*sys.argv[1:6])
+    create_shadertext(
+        shader_dir=Path(sys.argv[1]),
+        shader_dir2=Path(sys.argv[2]),
+        output_header=Path(sys.argv[3]),
+        output_source=Path(sys.argv[4]),
+    )
     create_buildinfo(dirname(sys.argv[4]), dirname(dirname(sys.argv[1])))