jw-python/tools/python/jwutils/stree/StringTree.py

from __future__ import annotations
from collections import OrderedDict
from typing import Optional, Union
from jwutils.log import *

def quote(s):
    if is_quoted(s):
        return s
    s = s.strip()
    if len(s) > 0:
        if s[0] == '"':
            return "'" + s + "'"
    return '"' + s + '"'

def is_quoted(s: str) -> bool:
    if isinstance(s, StringTree):
        return False
    s = s.strip()
    if len(s) < 2:
        return False
    d = s[0]
    if d == s[-1] and d in [ '"', "'" ]:
        return True
    return False

def cleanup_string(s: str) -> str:
    if isinstance(s, StringTree):
        return s
    s = s.strip()
    if is_quoted(s):
        return s[1:-1].replace('\\' + s[0], s[0])
    return s

class StringTree: # export

    def __init__(self, path: str, content: str) -> None:
        slog(DEBUG, "ctor, path =", path, "content =", content)
        self.children: OrderedDict[str, StringTree] = OrderedDict()
        self.content = None
        self.__set(path, content)
        assert(hasattr(self, "content"))
        #assert self.content is not None

    # root (content = [ symbols ])
    #  symbols (content = [ regex ])
    #    regex ( content ='[ \n\t\r]+' )

    # root (content = root, children = [ symbols ])
    #  symbols (content = symbols, children = [ regex ])
    #    regex ( content = regex, children = [ '[ \n\t\r]+' ] )
    #      '[ \n\t\r]+)' ( content = '\n\t\r]+)', children = [] )

    def __add_children(self, rhs):
        assert isinstance(rhs, StringTree)
        slog(DEBUG, "{}: adding children of {}: ".format(str(self), str(rhs)), str(rhs.children))
        #rhs.dump(INFO, "These children are added")
        self.content = rhs.content
        for name, c in rhs.children.items():
            if not name in self.children.keys():
                slog(DEBUG, "{}: adding new child: {}".format(str(self), str(c)))
                self.children[name] = c
            else:
                self.children[name].__add_children(c)

    def __set(self, path_, content, split=True):
        slog(DEBUG, "+ setting >" + str(content) + "< at path \"" + str(path_) + "\"" +
            (', containing "{}"'.format(str(self.content))) if hasattr(self, "content") else "")
        assert self.content != str(content)
        if content is not None and not type(content) in [str, StringTree]:
            raise Exception("Tried to add content of unsupported type {}".format(type(content).__name__))
        if path_ is None:
            if isinstance(content, str):
                self.content = cleanup_string(content)
            elif isinstance(content, StringTree):
                self.__add_children(content)
            else:
                raise Exception("Tried to add content of unsupported type {}".format(type(content).__name__))
            slog(DEBUG, " -- content = >" + str(content) + "<, self.content = >" + str(self.content) + "<")
            return self
        path = cleanup_string(path_)
        components = path.split('.') if split else [ path ]
        l = len(components)
        if len(path) == 0 or l == 0:
            #assert self.content is None or (isinstance(content, StringTree) and content.content == self.content)
            if isinstance(content, StringTree):
                #assert isinstance(content, StringTree), "Type: " + type(content).__name__
                self.__add_children(content)
            else:
                #self.content = cleanup_string(content)
                self.content = content
                slog(DEBUG, " -- content = >" + str(content) + "<, self.content = >" + str(self.content) + "<")
                #assert(content != "'antlr_doesnt_understand_vertical_tab'")
                assert(content != '"[a-zA-Z0-9+_*/-]"')
                #self.children[content] = StringTree(None, content)
            return self

        #assert self.content is not None, "tried to set empty content to {}".format(path_)

        nibble = components[0]
        rest = '.'.join(components[1:])
        if nibble not in self.children:
            self.children[nibble] = StringTree('', content=nibble)
        if l > 1:
            assert len(rest) > 0
            return self.children[nibble].__set(rest, content=content)
        if content is not None:
            gc = content if isinstance(content, StringTree) else StringTree('', content=content)
            self.children[nibble].children[gc.content] = gc
        return self.children[nibble]

    def __str__(self) -> str:
        return 'st:"{}"'.format(self.content)

    def __getitem__(self, path: str) -> str:
        r = self.get(path)
        if r is None:
            raise KeyError(path)
        return r.value()

    def __setitem__(self, key, value):
        return self.__set(key, value)

    def __dump(self, prio, indent=0, **kwargs):
        caller = kwargs['caller'] if 'caller' in kwargs.keys() else get_caller_pos(1)
        slog(prio, '|' + (' ' * indent) + str(self.content), caller=caller)
        indent += 2
        for name, child in self.children.items():
            child.__dump(prio, indent=indent, caller=caller)

    def keys(self):
        return self.children.keys()

    def items(self):
        return self.children.items()

    def set_content(self, content):
        if content is None:
            raise Exception("Tried to set none content")
        content = cleanup_string(content)
        if len(content) == 0:
            raise Exception("Tried to set empty content")
        self.content = content

    def add(self, path: str, content: Optional[Union[str, StringTree]] = None, split: bool = True) -> StringTree:
        slog(DEBUG, "adding >{}< at >{}< to >{}<".format(content, path, self.content))
        return self.__set(path, content, split)

    def get(self, path_: str) -> Optional[StringTree]:
        slog(DEBUG, 'looking for "{}" in "{}"'.format(path_, self.content))
        assert not isinstance(path_, int)
        path = cleanup_string(path_)
        if len(path) == 0:
            slog(DEBUG, "returning myself")
            return self
        if is_quoted(path_):
            if not path in self.children.keys():
                return None
            return self.children[path]
        components = path.split('.')
        if len(components) == 0:
            return self
        name = cleanup_string(components[0])
        if not hasattr(self, "children"):
            return None
        if not name in self.children.keys():
            slog(DEBUG, "Name \"" + name + "\" is not in children of", self.content)
            for child in self.children:
                slog(DEBUG, "child = ", child)
            return None
        relpath = '.'.join(components[1:])
        return self.children[name].get(relpath)

    def value(self) -> str:
        if len(self.children) == 0:
            raise Exception('tried to get value from leave "{}"'.format(self.content))
        return self.children[next(reversed(self.children))].content # type: ignore

    def child_list(self, depth_first=True):
        if depth_first == False:
            raise Exception("tried to retrieve child list with breadth-first search, not yet implemented")
        r = []
        for c in self.children:
            r.append(c)
            r.extend(c.to_list())

    def dump(self, prio: int, *args, **kwargs) -> None:
        caller = kwargs['caller'] if 'caller' in kwargs.keys() else get_caller_pos(1)
        msg = ''
        if args is not None:
            msg = ' ' + ' '.join(args) + ' '
        slog(prio, ",------------" + msg + "----------- >", caller=caller)
        self.__dump(prio, indent=0, caller=caller)
        slog(prio, "`------------" + msg + "----------- <", caller=caller)