From 1b134113834e65ad5453a8aa1548c9d41a40bcb7 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Sun, 13 Jun 2021 19:59:11 +0200 Subject: [PATCH] update ruamel.yaml, adjust api --- mitmproxy/optmanager.py | 18 +++++++++--------- mitmproxy/tools/console/keymap.py | 7 ++++--- mitmproxy/tools/main.py | 2 +- setup.py | 2 +- test/mitmproxy/test_optmanager.py | 20 +++++++++++++++----- 5 files changed, 30 insertions(+), 19 deletions(-) diff --git a/mitmproxy/optmanager.py b/mitmproxy/optmanager.py index a8720f79a..190faabb9 100644 --- a/mitmproxy/optmanager.py +++ b/mitmproxy/optmanager.py @@ -434,7 +434,7 @@ class OptManager: raise ValueError("Unsupported option type: %s", o.typespec) -def dump_defaults(opts): +def dump_defaults(opts, out: typing.TextIO): """ Dumps an annotated file with all options. """ @@ -453,7 +453,7 @@ def dump_defaults(opts): txt = "\n".join(textwrap.wrap(txt)) s.yaml_set_comment_before_after_key(k, before="\n" + txt) - return ruamel.yaml.round_trip_dump(s) + return ruamel.yaml.YAML().dump(s, out) def dump_dicts(opts, keys: typing.List[str]=None): @@ -482,7 +482,8 @@ def parse(text): if not text: return {} try: - data = ruamel.yaml.load(text, ruamel.yaml.RoundTripLoader) + yaml = ruamel.yaml.YAML(typ='unsafe', pure=True) + data = yaml.load(text) except ruamel.yaml.error.YAMLError as v: if hasattr(v, "problem_mark"): snip = v.problem_mark.get_snippet() @@ -532,7 +533,7 @@ def load_paths(opts: OptManager, *paths: str) -> None: ) -def serialize(opts: OptManager, text: str, defaults: bool = False) -> str: +def serialize(opts: OptManager, file: typing.TextIO, text: str, defaults: bool = False) -> None: """ Performs a round-trip serialization. If text is not None, it is treated as a previous serialization that should be modified @@ -550,9 +551,8 @@ def serialize(opts: OptManager, text: str, defaults: bool = False) -> str: for k in list(data.keys()): if k not in opts._options: del data[k] - ret = ruamel.yaml.round_trip_dump(data) - assert ret - return ret + + ruamel.yaml.YAML().dump(data, file) def save(opts: OptManager, path: str, defaults: bool =False) -> None: @@ -572,6 +572,6 @@ def save(opts: OptManager, path: str, defaults: bool =False) -> None: ) else: data = "" - data = serialize(opts, data, defaults) + with open(path, "wt", encoding="utf8") as f: - f.write(data) + serialize(opts, f, data, defaults) diff --git a/mitmproxy/tools/console/keymap.py b/mitmproxy/tools/console/keymap.py index d15853157..1ba7c9bd1 100644 --- a/mitmproxy/tools/console/keymap.py +++ b/mitmproxy/tools/console/keymap.py @@ -2,6 +2,7 @@ import typing import os import ruamel.yaml +import ruamel.yaml.error from mitmproxy import command from mitmproxy.tools.console import commandexecutor @@ -217,8 +218,8 @@ class KeymapConfig: def parse(self, text): try: - data = ruamel.yaml.safe_load(text) - except ruamel.yaml.error.YAMLError as v: + data = ruamel.yaml.YAML(typ='safe', pure=True).load(text) + except ruamel.yaml.error.MarkedYAMLError as v: if hasattr(v, "problem_mark"): snip = v.problem_mark.get_snippet() raise KeyBindingError( @@ -230,7 +231,7 @@ class KeymapConfig: if not data: return [] if not isinstance(data, list): - raise KeyBindingError("Inalid keybinding config - expected a list of keys") + raise KeyBindingError("Invalid keybinding config - expected a list of keys") for k in data: unknown = k.keys() - keyAttrs.keys() diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index ba0e8bf8e..60923d5b8 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -85,7 +85,7 @@ def run( process_options(parser, opts, args) if args.options: - print(optmanager.dump_defaults(opts)) + optmanager.dump_defaults(opts, sys.stdout) sys.exit(0) if args.commands: master.commands.dump() diff --git a/setup.py b/setup.py index 92a6c34a8..10df7d30f 100644 --- a/setup.py +++ b/setup.py @@ -84,7 +84,7 @@ setup( "pyOpenSSL>=20.0,<20.1", "pyparsing>=2.4.2,<2.5", "pyperclip>=1.6.0,<1.9", - "ruamel.yaml>=0.16,<0.17", + "ruamel.yaml>=0.16,<0.17.9", "sortedcontainers>=2.3,<2.5", "tornado>=4.3,<7", "urwid>=2.1.1,<2.2", diff --git a/test/mitmproxy/test_optmanager.py b/test/mitmproxy/test_optmanager.py index 9d109c229..fe8b59746 100644 --- a/test/mitmproxy/test_optmanager.py +++ b/test/mitmproxy/test_optmanager.py @@ -1,4 +1,5 @@ import copy +import io import pytest import typing import argparse @@ -239,11 +240,16 @@ def test_items(): def test_serialize(): + def serialize(opts: optmanager.OptManager, text: str, defaults: bool = False) -> str: + buf = io.StringIO() + optmanager.serialize(opts, buf, text, defaults) + return buf.getvalue() + o = TD2() o.three = "set" - assert "dfour" in optmanager.serialize(o, None, defaults=True) + assert "dfour" in serialize(o, "", defaults=True) - data = optmanager.serialize(o, None) + data = serialize(o, "") assert "dfour" not in data o2 = TD2() @@ -254,7 +260,7 @@ def test_serialize(): t = """ unknown: foo """ - data = optmanager.serialize(o, t) + data = serialize(o, t) o2 = TD2() optmanager.load(o2, data) assert o2 == o @@ -280,7 +286,9 @@ def test_serialize(): def test_serialize_defaults(): o = options.Options() - assert optmanager.serialize(o, None, defaults=True) + buf = io.StringIO() + optmanager.serialize(o, buf, "", defaults=True) + assert buf.getvalue() def test_saving(tmpdir): @@ -348,7 +356,9 @@ def test_option(): def test_dump_defaults(): o = TTypes() - assert optmanager.dump_defaults(o) + buf = io.StringIO() + optmanager.dump_defaults(o, buf) + assert buf.getvalue() def test_dump_dicts():