diff --git a/keynote_parser/file_utils.py b/keynote_parser/file_utils.py index db1f981..a4a48ab 100644 --- a/keynote_parser/file_utils.py +++ b/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,6 +161,8 @@ 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 @@ -121,22 +170,16 @@ def accept(filename, contents): 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 diff --git a/tests/test_performance_paths.py b/tests/test_performance_paths.py new file mode 100644 index 0000000..bf5bf85 --- /dev/null +++ b/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