tools/tryselect/preset.py

   1 # This Source Code Form is subject to the terms of the Mozilla Public
   2 # License, v. 2.0. If a copy of the MPL was not distributed with this
   3 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
   4
   5
   6 import os
   7 import shlex
   8 import subprocess
   9
  10 import yaml
  11
  12
  13 # Ensure strings with ' like fzf query strings get double-quoted:
  14 def represent_str(self, data):
  15     if "'" in data:
  16         return self.represent_scalar("tag:yaml.org,2002:str", data, style='"')
  17     return self.represent_scalar("tag:yaml.org,2002:str", data)
  18
  19
  20 yaml.SafeDumper.add_representer(str, represent_str)
  21
  22
  23 class PresetHandler:
  24     def __init__(self, path):
  25         self.path = path
  26         self._presets = {}
  27
  28     @property
  29     def presets(self):
  30         if not self._presets and os.path.isfile(self.path):
  31             with open(self.path) as fh:
  32                 self._presets = yaml.safe_load(fh) or {}
  33
  34         return self._presets
  35
  36     def __contains__(self, name):
  37         return name in self.presets
  38
  39     def __getitem__(self, name):
  40         return self.presets[name]
  41
  42     def __len__(self):
  43         return len(self.presets)
  44
  45     def __str__(self):
  46         if not self.presets:
  47             return ""
  48         return yaml.safe_dump(self.presets, default_flow_style=False)
  49
  50     def list(self):
  51         if not self.presets:
  52             print("no presets found")
  53         else:
  54             print(self)
  55
  56     def edit(self):
  57         if "EDITOR" not in os.environ:
  58             print(
  59                 "error: must set the $EDITOR environment variable to use --edit-presets"
  60             )
  61             return
  62
  63         subprocess.call(shlex.split(os.environ["EDITOR"]) + [self.path])
  64
  65     def save(self, name, **data):
  66         self.presets[name] = data
  67
  68         with open(self.path, "w") as fh:
  69             fh.write(str(self))
  70
  71
  72 class MergedHandler:
  73     def __init__(self, *paths):
  74         """Helper class for dealing with multiple preset files."""
  75         self.handlers = [PresetHandler(p) for p in paths]
  76
  77     def __contains__(self, name):
  78         return any(name in handler for handler in self.handlers)
  79
  80     def __getitem__(self, name):
  81         for handler in self.handlers:
  82             if name in handler:
  83                 return handler[name]
  84         raise KeyError(name)
  85
  86     def __len__(self):
  87         return sum(len(h) for h in self.handlers)
  88
  89     def __str__(self):
  90         all_presets = {
  91             k: v for handler in self.handlers for k, v in handler.presets.items()
  92         }
  93         return yaml.safe_dump(all_presets, default_flow_style=False)
  94
  95     def list(self):
  96         if len(self) == 0:
  97             print("no presets found")
  98             return
  99
 100         for handler in self.handlers:
 101             val = str(handler)
 102             if val:
 103                 val = "\n  ".join(
 104                     [""] + val.splitlines() + [""]
 105                 )  # indent all lines by 2 spaces
 106                 print("Presets from {}:".format(handler.path))
 107                 print(val)