Stop doing work the output doesn't depend on

keynote_parser/file_utils.py

@@ -23,6 +23,56 @@
 from .unicode_utils import fix_unicode
 
 
+# Resolver.resolve() decides a scalar's implicit tag by running the resolver
+# table's regexes against its value. It is a pure function of (value, implicit)
+# for a fixed table, and PyYAML calls it once per scalar - 259,176 times for a
+# single 320 KB .key file, because the YAML form of a document is ~30x the size
+# of the .iwa it came from. Memoizing it produces byte-identical output.
+#
+# CDumper only replaces the emitter; the Representer and Resolver above it stay
+# pure Python, which is why YAML dominates the profile even with libyaml
+# installed.
+_RESOLVE_CACHE_MAX_ENTRIES = 100_000
+
+
+class MemoizingDumper(Dumper):
+    _resolve_cache = {}
+
+    def resolve(self, kind, value, implicit):
+        if kind is not yaml.ScalarNode:
+            return super().resolve(kind, value, implicit)
+        key = (value, implicit)
+        cached = self._resolve_cache.get(key)
+        if cached is not None:
+            return cached
+        resolved = super().resolve(kind, value, implicit)
+        # Bounded so a very large document can't grow this without limit.
+        if len(self._resolve_cache) < _RESOLVE_CACHE_MAX_ENTRIES:
+            self._resolve_cache[key] = resolved
+        return resolved
+
+
+def dump_yaml(data, stream=None):
+    """Serialize an IWAFile's dict form to YAML, the one way we do it."""
+    return yaml.dump(
+        data,
+        stream,
+        default_flow_style=False,
+        encoding="utf-8",
+        Dumper=MemoizingDumper,
+    )
+
+
+def to_yaml_data(contents):
+    """Accept either an IWAFile or a dict that has already been built from one.
+
+    process_file hands the sink a plain dict when a replacement has already
+    produced one, so that writing YAML doesn't have to round-trip it back
+    through IWAFile.from_dict() and then to_dict() again.
+    """
+    return contents.to_dict() if isinstance(contents, IWAFile) else contents
+
+
 def ensure_directory_exists(prefix, path):
     """Ensure that a path's directory exists."""
     parts = os.path.split(path)
@@ -87,24 +137,21 @@ def dir_file_sink(target_dir, raw=False):
     def accept(filename, contents):
         ensure_directory_exists(target_dir, filename)
         target_path = os.path.join(target_dir, filename)
-        if isinstance(contents, IWAFile) and not raw:
+        is_iwa = isinstance(contents, (IWAFile, dict))
+        if is_iwa and not raw:
             target_path += ".yaml"
         with open(target_path, "wb") as out:
-            if isinstance(contents, IWAFile):
+            if is_iwa:
                 if raw:
                     out.write(contents.to_buffer())
                 else:
-                    yaml.dump(
-                        contents.to_dict(),
-                        out,
-                        default_flow_style=False,
-                        encoding="utf-8",
-                        Dumper=Dumper,
-                    )
+                    dump_yaml(to_yaml_data(contents), out)
             else:
                 out.write(contents)
 
     accept.uses_stdout = False
+    # Writing YAML needs the dict, not a rebuilt IWAFile.
+    accept.needs_iwa_object = raw
     yield accept
 
 
@@ -114,29 +161,25 @@ def accept(filename, contents):
         print(filename)
 
     accept.uses_stdout = True
+    # `ls` only ever prints names, so there is no reason to decode anything.
+    accept.needs_iwa_contents = False
     yield accept
 
 
 @contextmanager
 def cat_sink(subfile, raw):
     def accept(filename, contents):
         if filename == subfile:
-            if isinstance(contents, IWAFile):
+            if isinstance(contents, (IWAFile, dict)):
                 if raw:
                     sys.stdout.buffer.write(contents.to_buffer())
                 else:
-                    print(
-                        yaml.dump(
-                            contents.to_dict(),
-                            default_flow_style=False,
-                            encoding="utf-8",
-                            Dumper=Dumper,
-                        ).decode("ascii")
-                    )
+                    print(dump_yaml(to_yaml_data(contents)).decode("ascii"))
             else:
                 sys.stdout.buffer.write(contents)
 
     accept.uses_stdout = True
+    accept.needs_iwa_object = raw
     yield accept
 
 
@@ -148,6 +191,8 @@ def accept(filename, contents):
         files_to_write[filename] = contents
 
     accept.uses_stdout = False
+    # Writes binary .iwa back out, so it needs a real IWAFile to serialize.
+    accept.needs_iwa_object = True
 
     yield accept
 
@@ -165,6 +210,12 @@ def accept(filename, contents):
 def process_file(filename, handle, sink, replacements=[], raw=False, on_replace=None):
     contents = None
     if ".iwa" in filename and not raw:
+        if not getattr(sink, "needs_iwa_contents", True):
+            # `ls` only prints names. Decoding every archive to throw the result
+            # away made listing a document cost as much as parsing it.
+            sink(filename.replace(".yaml", ""), None)
+            return
+
         contents = handle.read()
         if filename.endswith(".yaml"):
             file = IWAFile.from_dict(
@@ -186,7 +237,13 @@ def process_file(filename, handle, sink, replacements=[], raw=False, on_replace=
             data = file.to_dict()
             for replacement in replacements:
                 data = replacement.perform_on(data, on_replace=on_replace)
-            sink(filename, IWAFile.from_dict(data))
+            if getattr(sink, "needs_iwa_object", True):
+                sink(filename, IWAFile.from_dict(data))
+            else:
+                # A YAML sink is about to call to_dict() on whatever it is
+                # given, so rebuilding an IWAFile from this dict only to take
+                # it apart again is pure overhead.
+                sink(filename, data)
         else:
             sink(filename, file)
         return

tests/test_performance_paths.py

@@ -0,0 +1,122 @@
+"""Tests for the shortcuts taken on the hot paths.
+
+These are correctness tests, not benchmarks: each optimisation skips work, and
+what matters is that skipping it changes nothing observable.
+"""
+
+import io
+import zipfile
+
+import pytest
+import yaml
+
+from keynote_parser import file_utils
+from keynote_parser.codec import IWAFile
+from keynote_parser.replacement import Replacement
+
+TABLE_FILENAME = "./tests/data/table.key"
+
+
+def _iwa_names(path):
+    with zipfile.ZipFile(path) as archive:
+        return [n for n in archive.namelist() if n.endswith(".iwa")]
+
+
+def test_ls_does_not_decode_archives(monkeypatch):
+    """`ls` prints names, so it must not pay to parse every archive."""
+    calls = []
+    original = IWAFile.from_buffer
+
+    def counting_from_buffer(*args, **kwargs):
+        calls.append(args[1] if len(args) > 1 else None)
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(IWAFile, "from_buffer", counting_from_buffer)
+
+    with file_utils.file_sink("-") as sink:
+        for name, handle in file_utils.file_reader(TABLE_FILENAME, False):
+            file_utils.process_file(name, handle, sink)
+
+    assert calls == [], f"ls decoded {len(calls)} archives it never looked at"
+
+
+def test_ls_still_lists_every_file(capsys):
+    with file_utils.file_sink("-") as sink:
+        for name, handle in file_utils.file_reader(TABLE_FILENAME, False):
+            file_utils.process_file(name, handle, sink)
+    printed = set(capsys.readouterr().out.split())
+
+    with zipfile.ZipFile(TABLE_FILENAME) as archive:
+        expected = {n for n in archive.namelist() if not n.endswith("/")}
+    assert printed == expected
+
+
+def test_ls_of_a_directory_strips_the_yaml_suffix(tmp_path, capsys):
+    """Unpacked directories hold .iwa.yaml; `ls` has always reported .iwa."""
+    out = str(tmp_path / "unpacked")
+    file_utils.process(TABLE_FILENAME, out)
+    capsys.readouterr()
+
+    with file_utils.file_sink("-") as sink:
+        for name, handle in file_utils.file_reader(out, False):
+            file_utils.process_file(name, handle, sink)
+    printed = set(capsys.readouterr().out.split())
+
+    assert any(n.endswith(".iwa") for n in printed)
+    assert not any(n.endswith(".yaml") for n in printed)
+
+
+@pytest.mark.parametrize("with_replacement", [False, True])
+def test_yaml_output_is_unchanged_by_the_shortcuts(tmp_path, with_replacement):
+    """The dict handed to a YAML sink must serialize exactly as an IWAFile would."""
+    name = _iwa_names(TABLE_FILENAME)[0]
+    with zipfile.ZipFile(TABLE_FILENAME) as archive:
+        file = IWAFile.from_buffer(archive.read(name), name)
+
+    data = file.to_dict()
+    if with_replacement:
+        data = Replacement("e", "E").perform_on(data)
+
+    # What a YAML sink now writes, versus rebuilding an IWAFile first.
+    direct = file_utils.dump_yaml(file_utils.to_yaml_data(data))
+    roundtripped = file_utils.dump_yaml(
+        file_utils.to_yaml_data(IWAFile.from_dict(data))
+    )
+    assert direct == roundtripped
+
+
+def test_memoizing_dumper_matches_the_plain_one():
+    with zipfile.ZipFile(TABLE_FILENAME) as archive:
+        names = _iwa_names(TABLE_FILENAME)[:6]
+        dicts = [IWAFile.from_buffer(archive.read(n), n).to_dict() for n in names]
+
+    for data in dicts:
+        expected = yaml.dump(
+            data,
+            default_flow_style=False,
+            encoding="utf-8",
+            Dumper=file_utils.Dumper,
+        )
+        assert file_utils.dump_yaml(data) == expected
+
+
+def test_memoizing_dumper_resolves_ambiguous_scalars_correctly():
+    """Values that look like other types must keep their quoting."""
+    tricky = {
+        "digits": "123",
+        "float_ish": "1.5",
+        "bool_ish": ["yes", "no", "true", "false", "on", "off"],
+        "null_ish": ["null", "~", ""],
+        "real_int": 123,
+        "real_bool": True,
+        "real_none": None,
+        "sexagesimal": "1:30",
+        "leading_zero": "007",
+    }
+    memoized = file_utils.dump_yaml(tricky)
+    plain = yaml.dump(
+        tricky, default_flow_style=False, encoding="utf-8", Dumper=file_utils.Dumper
+    )
+    assert memoized == plain
+    # And it must survive a round trip unchanged.
+    assert yaml.load(io.BytesIO(memoized), Loader=file_utils.Loader) == tricky