codekingpro/portable-devtools

1from __future__ import annotations2 3import contextlib4import copy5import pprint6import textwrap7import weakref8from collections.abc import Callable9from collections.abc import Iterable10from collections.abc import Sequence11from dataclasses import dataclass12from pathlib import Path13from typing import Any14from typing import Optional15from typing import TextIO16 17import ruamel.yaml18 19from mitmproxy import exceptions20from mitmproxy.utils import signals21from mitmproxy.utils import typecheck22 23"""24    The base implementation for Options.25"""26 27unset = object()28 29 30class _Option:31    __slots__ = ("name", "typespec", "value", "_default", "choices", "help")32 33    def __init__(34        self,35        name: str,36        typespec: type | object,  # object for Optional[x], which is not a type.37        default: Any,38        help: str,39        choices: Sequence[str] | None,40    ) -> None:41        typecheck.check_option_type(name, default, typespec)42        self.name = name43        self.typespec = typespec44        self._default = default45        self.value = unset46        self.help = textwrap.dedent(help).strip().replace("\n", " ")47        self.choices = choices48 49    def __repr__(self):50        return f"{self.current()} [{self.typespec}]"51 52    @property53    def default(self):54        return copy.deepcopy(self._default)55 56    def current(self) -> Any:57        if self.value is unset:58            v = self.default59        else:60            v = self.value61        return copy.deepcopy(v)62 63    def set(self, value: Any) -> None:64        typecheck.check_option_type(self.name, value, self.typespec)65        self.value = value66 67    def reset(self) -> None:68        self.value = unset69 70    def has_changed(self) -> bool:71        return self.current() != self.default72 73    def __eq__(self, other) -> bool:74        for i in self.__slots__:75            if getattr(self, i) != getattr(other, i):76                return False77        return True78 79    def __deepcopy__(self, _):80        o = _Option(self.name, self.typespec, self.default, self.help, self.choices)81        if self.has_changed():82            o.value = self.current()83        return o84 85 86@dataclass87class _UnconvertedStrings:88    val: list[str]89 90 91def _sig_changed_spec(updated: set[str]) -> None:  # pragma: no cover92    ...  # expected function signature for OptManager.changed receivers.93 94 95def _sig_errored_spec(exc: Exception) -> None:  # pragma: no cover96    ...  # expected function signature for OptManager.errored receivers.97 98 99class OptManager:100    """101    OptManager is the base class from which Options objects are derived.102 103    .changed is a Signal that triggers whenever options are104    updated. If any handler in the chain raises an exceptions.OptionsError105    exception, all changes are rolled back, the exception is suppressed,106    and the .errored signal is notified.107 108    Optmanager always returns a deep copy of options to ensure that109    mutation doesn't change the option state inadvertently.110    """111 112    def __init__(self) -> None:113        self.deferred: dict[str, Any] = {}114        self.changed = signals.SyncSignal(_sig_changed_spec)115        self.changed.connect(self._notify_subscribers)116        self.errored = signals.SyncSignal(_sig_errored_spec)117        self._subscriptions: list[tuple[weakref.ref[Callable], set[str]]] = []118        # Options must be the last attribute here - after that, we raise an119        # error for attribute assignment to unknown options.120        self._options: dict[str, Any] = {}121 122    def add_option(123        self,124        name: str,125        typespec: type | object,126        default: Any,127        help: str,128        choices: Sequence[str] | None = None,129    ) -> None:130        self._options[name] = _Option(name, typespec, default, help, choices)131        self.changed.send(updated={name})132 133    @contextlib.contextmanager134    def rollback(self, updated, reraise=False):135        old = copy.deepcopy(self._options)136        try:137            yield138        except exceptions.OptionsError as e:139            # Notify error handlers140            self.errored.send(exc=e)141            # Rollback142            self.__dict__["_options"] = old143            self.changed.send(updated=updated)144            if reraise:145                raise e146 147    def subscribe(self, func, opts):148        """149        Subscribe a callable to the .changed signal, but only for a150        specified list of options. The callable should accept arguments151        (options, updated), and may raise an OptionsError.152 153        The event will automatically be unsubscribed if the callable goes out of scope.154        """155        for i in opts:156            if i not in self._options:157                raise exceptions.OptionsError("No such option: %s" % i)158 159        self._subscriptions.append((signals.make_weak_ref(func), set(opts)))160 161    def _notify_subscribers(self, updated) -> None:162        cleanup = False163        for ref, opts in self._subscriptions:164            callback = ref()165            if callback is not None:166                if opts & updated:167                    callback(self, updated)168            else:169                cleanup = True170 171        if cleanup:172            self.__dict__["_subscriptions"] = [173                (ref, opts) for (ref, opts) in self._subscriptions if ref() is not None174            ]175 176    def __eq__(self, other):177        if isinstance(other, OptManager):178            return self._options == other._options179        return False180 181    def __deepcopy__(self, memodict=None):182        o = OptManager()183        o.__dict__["_options"] = copy.deepcopy(self._options, memodict)184        return o185 186    __copy__ = __deepcopy__187 188    def __getattr__(self, attr):189        if attr in self._options:190            return self._options[attr].current()191        else:192            raise AttributeError("No such option: %s" % attr)193 194    def __setattr__(self, attr, value):195        # This is slightly tricky. We allow attributes to be set on the instance196        # until we have an _options attribute. After that, assignment is sent to197        # the update function, and will raise an error for unknown options.198        opts = self.__dict__.get("_options")199        if not opts:200            super().__setattr__(attr, value)201        else:202            self.update(**{attr: value})203 204    def keys(self):205        return set(self._options.keys())206 207    def items(self):208        return self._options.items()209 210    def __contains__(self, k):211        return k in self._options212 213    def reset(self):214        """215        Restore defaults for all options.216        """217        for o in self._options.values():218            o.reset()219        self.changed.send(updated=set(self._options.keys()))220 221    def update_known(self, **kwargs):222        """223        Update and set all known options from kwargs. Returns a dictionary224        of unknown options.225        """226        known, unknown = {}, {}227        for k, v in kwargs.items():228            if k in self._options:229                known[k] = v230            else:231                unknown[k] = v232        updated = set(known.keys())233        if updated:234            with self.rollback(updated, reraise=True):235                for k, v in known.items():236                    self._options[k].set(v)237                self.changed.send(updated=updated)238        return unknown239 240    def update_defer(self, **kwargs):241        unknown = self.update_known(**kwargs)242        self.deferred.update(unknown)243 244    def update(self, **kwargs):245        u = self.update_known(**kwargs)246        if u:247            raise KeyError("Unknown options: %s" % ", ".join(u.keys()))248 249    def setter(self, attr):250        """251        Generate a setter for a given attribute. This returns a callable252        taking a single argument.253        """254        if attr not in self._options:255            raise KeyError("No such option: %s" % attr)256 257        def setter(x):258            setattr(self, attr, x)259 260        return setter261 262    def toggler(self, attr):263        """264        Generate a toggler for a boolean attribute. This returns a callable265        that takes no arguments.266        """267        if attr not in self._options:268            raise KeyError("No such option: %s" % attr)269        o = self._options[attr]270        if o.typespec is not bool:271            raise ValueError("Toggler can only be used with boolean options")272 273        def toggle():274            setattr(self, attr, not getattr(self, attr))275 276        return toggle277 278    def default(self, option: str) -> Any:279        return self._options[option].default280 281    def has_changed(self, option):282        """283        Has the option changed from the default?284        """285        return self._options[option].has_changed()286 287    def merge(self, opts):288        """289        Merge a dict of options into this object. Options that have None290        value are ignored. Lists and tuples are appended to the current291        option value.292        """293        toset = {}294        for k, v in opts.items():295            if v is not None:296                if isinstance(v, (list, tuple)):297                    toset[k] = getattr(self, k) + v298                else:299                    toset[k] = v300        self.update(**toset)301 302    def __repr__(self):303        options = pprint.pformat(self._options, indent=4).strip(" {}")304        if "\n" in options:305            options = "\n    " + options + "\n"306        return "{mod}.{cls}({{{options}}})".format(307            mod=type(self).__module__, cls=type(self).__name__, options=options308        )309 310    def set(self, *specs: str, defer: bool = False) -> None:311        """312        Takes a list of set specification in standard form (option=value).313        Options that are known are updated immediately. If defer is true,314        options that are not known are deferred, and will be set once they315        are added.316 317        May raise an `OptionsError` if a value is malformed or an option is unknown and defer is False.318        """319        # First, group specs by option name.320        unprocessed: dict[str, list[str]] = {}321        for spec in specs:322            if "=" in spec:323                name, value = spec.split("=", maxsplit=1)324                unprocessed.setdefault(name, []).append(value)325            else:326                unprocessed.setdefault(spec, [])327 328        # Second, convert values to the correct type.329        processed: dict[str, Any] = {}330        for name in list(unprocessed.keys()):331            if name in self._options:332                processed[name] = self._parse_setval(333                    self._options[name], unprocessed.pop(name)334                )335 336        # Third, stash away unrecognized options or complain about them.337        if defer:338            self.deferred.update(339                {k: _UnconvertedStrings(v) for k, v in unprocessed.items()}340            )341        elif unprocessed:342            raise exceptions.OptionsError(343                f"Unknown option(s): {', '.join(unprocessed)}"344            )345 346        # Finally, apply updated options.347        self.update(**processed)348 349    def process_deferred(self) -> None:350        """351        Processes options that were deferred in previous calls to set, and352        have since been added.353        """354        update: dict[str, Any] = {}355        for optname, value in self.deferred.items():356            if optname in self._options:357                if isinstance(value, _UnconvertedStrings):358                    value = self._parse_setval(self._options[optname], value.val)359                update[optname] = value360        self.update(**update)361        for k in update.keys():362            del self.deferred[k]363 364    def _parse_setval(self, o: _Option, values: list[str]) -> Any:365        """366        Convert a string to a value appropriate for the option type.367        """368        if o.typespec == Sequence[str]:369            return values370        if len(values) > 1:371            raise exceptions.OptionsError(372                f"Received multiple values for {o.name}: {values}"373            )374 375        optstr: str | None376        if values:377            optstr = values[0]378        else:379            optstr = None380 381        if o.typespec in (str, Optional[str]):382            if o.typespec is str and optstr is None:383                raise exceptions.OptionsError(f"Option is required: {o.name}")384            return optstr385        elif o.typespec in (int, Optional[int]):386            if optstr:387                try:388                    return int(optstr)389                except ValueError:390                    raise exceptions.OptionsError(391                        f"Failed to parse option {o.name}: not an integer: {optstr}"392                    )393            elif o.typespec is int:394                raise exceptions.OptionsError(f"Option is required: {o.name}")395            else:396                return None397        elif o.typespec is bool:398            if optstr == "toggle":399                return not o.current()400            if not optstr or optstr == "true":401                return True402            elif optstr == "false":403                return False404            else:405                raise exceptions.OptionsError(406                    f'Failed to parse option {o.name}: boolean must be "true", "false", or have the value omitted (a synonym for "true").'407                )408        raise NotImplementedError(409            f"Failed to parse option {o.name}: unsupported option type: {o.typespec}"410        )411 412    def make_parser(self, parser, optname, metavar=None, short=None):413        """414        Auto-Create a command-line parser entry for a named option. If the415        option does not exist, it is ignored.416        """417        if optname not in self._options:418            return419 420        o = self._options[optname]421 422        def mkf(x, s):423            x = x.replace("_", "-")424            f = ["--%s" % x]425            if s:426                f.append("-" + s)427            return f428 429        flags = mkf(optname, short)430 431        if o.typespec is bool:432            g = parser.add_mutually_exclusive_group(required=False)433            onf = mkf(optname, None)434            offf = mkf("no-" + optname, None)435            # The short option for a bool goes to whatever is NOT the default436            if short:437                if o.default:438                    offf = mkf("no-" + optname, short)439                else:440                    onf = mkf(optname, short)441            g.add_argument(442                *offf,443                action="store_false",444                dest=optname,445            )446            g.add_argument(*onf, action="store_true", dest=optname, help=o.help)447            parser.set_defaults(**{optname: None})448        elif o.typespec in (int, Optional[int]):449            parser.add_argument(450                *flags,451                action="store",452                type=int,453                dest=optname,454                help=o.help,455                metavar=metavar,456            )457        elif o.typespec in (str, Optional[str]):458            parser.add_argument(459                *flags,460                action="store",461                type=str,462                dest=optname,463                help=o.help,464                metavar=metavar,465                choices=o.choices,466            )467        elif o.typespec == Sequence[str]:468            parser.add_argument(469                *flags,470                action="append",471                type=str,472                dest=optname,473                help=o.help + " May be passed multiple times.",474                metavar=metavar,475                choices=o.choices,476            )477        else:478            raise ValueError("Unsupported option type: %s", o.typespec)479 480 481def dump_defaults(opts, out: TextIO):482    """483    Dumps an annotated file with all options.484    """485    # Sort data486    s = ruamel.yaml.comments.CommentedMap()487    for k in sorted(opts.keys()):488        o = opts._options[k]489        s[k] = o.default490        txt = o.help.strip()491 492        if o.choices:493            txt += " Valid values are %s." % ", ".join(repr(c) for c in o.choices)494        else:495            t = typecheck.typespec_to_str(o.typespec)496            txt += " Type %s." % t497 498        txt = "\n".join(textwrap.wrap(txt))499        s.yaml_set_comment_before_after_key(k, before="\n" + txt)500    return ruamel.yaml.YAML().dump(s, out)501 502 503def dump_dicts(opts, keys: Iterable[str] | None = None) -> dict:504    """505    Dumps the options into a list of dict object.506 507    Return: A list like: { "anticache": { type: "bool", default: false, value: true, help: "help text"} }508    """509    options_dict = {}510    if keys is None:511        keys = opts.keys()512    for k in sorted(keys):513        o = opts._options[k]514        t = typecheck.typespec_to_str(o.typespec)515        option = {516            "type": t,517            "default": o.default,518            "value": o.current(),519            "help": o.help,520            "choices": o.choices,521        }522        options_dict[k] = option523    return options_dict524 525 526def parse(text):527    if not text:528        return {}529    try:530        yaml = ruamel.yaml.YAML(typ="safe", pure=True)531        data = yaml.load(text)532    except ruamel.yaml.error.YAMLError as v:533        if hasattr(v, "problem_mark"):534            snip = v.problem_mark.get_snippet()535            raise exceptions.OptionsError(536                "Config error at line %s:\n%s\n%s"537                % (v.problem_mark.line + 1, snip, getattr(v, "problem", ""))538            )539        else:540            raise exceptions.OptionsError("Could not parse options.")541    if isinstance(data, str):542        raise exceptions.OptionsError("Config error - no keys found.")543    elif data is None:544        return {}545    return data546 547 548def load(opts: OptManager, text: str, cwd: Path | str | None = None) -> None:549    """550    Load configuration from text, over-writing options already set in551    this object. May raise OptionsError if the config file is invalid.552    """553    data = parse(text)554 555    scripts = data.get("scripts")556    if scripts is not None and cwd is not None:557        data["scripts"] = [558            str(relative_path(Path(path), relative_to=Path(cwd))) for path in scripts559        ]560 561    opts.update_defer(**data)562 563 564def load_paths(opts: OptManager, *paths: Path | str) -> None:565    """566    Load paths in order. Each path takes precedence over the previous567    path. Paths that don't exist are ignored, errors raise an568    OptionsError.569    """570    for p in paths:571        p = Path(p).expanduser()572        if p.exists() and p.is_file():573            with p.open(encoding="utf8") as f:574                try:575                    txt = f.read()576                except UnicodeDecodeError as e:577                    raise exceptions.OptionsError(f"Error reading {p}: {e}")578            try:579                load(opts, txt, cwd=p.absolute().parent)580            except exceptions.OptionsError as e:581                raise exceptions.OptionsError(f"Error reading {p}: {e}")582 583 584def serialize(585    opts: OptManager, file: TextIO, text: str, defaults: bool = False586) -> None:587    """588    Performs a round-trip serialization. If text is not None, it is589    treated as a previous serialization that should be modified590    in-place.591 592    - If "defaults" is False, only options with non-default values are593        serialized. Default values in text are preserved.594    - Unknown options in text are removed.595    - Raises OptionsError if text is invalid.596    """597    data = parse(text)598    for k in opts.keys():599        if defaults or opts.has_changed(k):600            data[k] = getattr(opts, k)601    for k in list(data.keys()):602        if k not in opts._options:603            del data[k]604 605    ruamel.yaml.YAML().dump(data, file)606 607 608def save(opts: OptManager, path: Path | str, defaults: bool = False) -> None:609    """610    Save to path. If the destination file exists, modify it in-place.611 612    Raises OptionsError if the existing data is corrupt.613    """614    path = Path(path).expanduser()615    if path.exists() and path.is_file():616        with path.open(encoding="utf8") as f:617            try:618                data = f.read()619            except UnicodeDecodeError as e:620                raise exceptions.OptionsError(f"Error trying to modify {path}: {e}")621    else:622        data = ""623 624    with path.open("w", encoding="utf8") as f:625        serialize(opts, f, data, defaults)626 627 628def relative_path(script_path: Path | str, *, relative_to: Path | str) -> Path:629    """630    Make relative paths found in config files relative to said config file,631    instead of relative to where the command is ran.632    """633    script_path = Path(script_path)634    # Edge case when $HOME is not an absolute path635    if script_path.expanduser() != script_path and not script_path.is_absolute():636        script_path = script_path.expanduser().absolute()637    return (relative_to / script_path.expanduser()).absolute()638