Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
YDefsLoader: add variable interpolation support
Browse files Browse the repository at this point in the history
the default YAML format does not support variable interpolation
e.g. referring to another key's value, but it's possible to add
it with the help of 'add_constructor` and `add_implicit_resolver`.

Example usage:

```yaml
c: cool
x:
  y: abcdefg
  z: hijklmno
  t: ${c}

foo: ${x.y}${x.z} ${x.t}
```

Signed-off-by: Mustafa Kemal Gilor <mustafa.gilor@canonical.com>
xmkg committed May 8, 2024
1 parent 4b97ef1 commit bfedd2a
Showing 1 changed file with 106 additions and 2 deletions.
108 changes: 106 additions & 2 deletions hotsos/core/ycheck/engine/common.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,113 @@
import abc
import os
import re

import yaml
from hotsos.core.config import HotSOSConfig
from hotsos.core.log import log


class YRefKeyNotFoundException(Exception):
def __init__(self, key):
message = f"{key} could not be found."
super().__init__(message)


class YRefReachedMaxAttemptsException(Exception):
def __init__(self, key):
message = f"Max search attempts have been reached for {key}"
super().__init__(message)


class YRefNotAScalarValueException(Exception):
def __init__(self, key, type_name):
message = f"{key} has non-scalar value type ({type_name})"
super().__init__(message)


class YSafeRefLoader(yaml.SafeLoader):
"""This class is just the regular yaml.SafeLoader but also resolves the
variable names to their values in YAML, e.g.;
x:
y: abc
z: def
foo : ${x.y} ${x.z}
# foo's value would be "abc def"
"""

# The regex pattern for detecting the variable names.
ref_matcher = None

def __init__(self, stream):
super().__init__(stream)
if not YSafeRefLoader.ref_matcher:
YSafeRefLoader.ref_matcher = re.compile(r'\$\{([^}^{]+)\}')
# register a custom tag for which our constructor is called
YSafeRefLoader.add_constructor("!ref",
YSafeRefLoader.ref_constructor)

# tell PyYAML that a scalar that looks like `${...}` is to be
# implicitly tagged with `!ref`, so that our custom constructor
# is called.
YSafeRefLoader.add_implicit_resolver("!ref",
YSafeRefLoader.ref_matcher,
None)

# we override this method to remember the root node,
# so that we can later resolve paths relative to it
def get_single_node(self):
self.cur_root = super().get_single_node()
return self.cur_root

@staticmethod
def ref_constructor(loader, node):
cur = loader.cur_root

max_resolve_attempts = 1000 # arbitrary choice

while max_resolve_attempts:
max_resolve_attempts -= 1
var = YSafeRefLoader.ref_matcher.search(node.value)
if not var:
break
target_key = var.group(1)
key_segments = target_key.split(".")
# Try to resolve the target variable
while key_segments:
# Get the segment on the front
current_segment = key_segments.pop(0)
found = False
# Iterate over current node's children
for (key, value) in cur.value:
# Check if node name matches with the current segment
if key.value == current_segment:
found = True
# we're the end of the segments, so we've
# reached to the node we want
if not key_segments:
ref_value = loader.construct_object(value)
if isinstance(ref_value, (dict, list)):
raise YRefNotAScalarValueException(
target_key,
type(ref_value))
node.value = node.value[:var.span()[0]] + \
ref_value + node.value[var.span()[1]:]
break
# Set the current node as root for key search
cur = value
break

if not found:
raise YRefKeyNotFoundException(target_key)

if not max_resolve_attempts:
raise YRefReachedMaxAttemptsException(target_key)

return node.value


class YDefsLoader(object):
""" Load yaml definitions. """

@@ -39,7 +141,8 @@ def _get_defs_recursive(self, path):
if self._get_yname(abs_path) == os.path.basename(path):
with open(abs_path) as fd:
log.debug("applying dir globals %s", entry)
defs.update(yaml.safe_load(fd.read()) or {})
defs.update(yaml.load(fd.read(),
Loader=YSafeRefLoader) or {})

# NOTE: these files do not count towards the total loaded
# since they are only supposed to contain directory-level
@@ -49,7 +152,8 @@ def _get_defs_recursive(self, path):

with open(abs_path) as fd:
self.stats_num_files_loaded += 1
_content = yaml.safe_load(fd.read()) or {}
_content = yaml.load(fd.read(),
Loader=YSafeRefLoader) or {}
defs[self._get_yname(abs_path)] = _content

return defs

0 comments on commit bfedd2a

Please sign in to comment.