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

from __future__ import annotations
from collections import OrderedDict
from typing import Any, List, 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, f'Constructing StringTree(path="{path}", content="{content}")')
        self.children: OrderedDict[str, StringTree] = OrderedDict()
        self.content: Optional[str] = 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 __adopt_children(self, parent):
        assert isinstance(parent, StringTree)
        slog(DEBUG, f'At {self.content}: Adopting children of {parent}')
        #parent.dump(INFO, "These children are added")
        self.content = parent.content
        for name, c in parent.children.items():
            if not name in self.children.keys():
                slog(DEBUG, f'At {self.content}: Adding new child {c}')
                self.children[name] = c
            else:
                self.children[name].__adopt_children(c)

    def __set(self, path_, content, split=True):
        slog(DEBUG, ('At "{}": '.format(str(self.content)) if hasattr(self, "content") else "") + f'Setting "{path_}" -> "{content}"')
        #assert self.content != str(content) # Not sure what the idea behind this was. It often goes off, and all works fine without.
        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.__adopt_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.__adopt_children(content)
            else:
                if self.content != content:
                    #self.content = cleanup_string(content)
                    slog(DEBUG, f'Changing content: "{self.content}" ->"{content}"')
                    assert(content != '"[a-zA-Z0-9+_*/-]"')
                    self.content = content
                    #assert(content != "'antlr_doesnt_understand_vertical_tab'")
                    #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)
        # last component, a.k.a. leaf
        if content is not None:
            gc = content if isinstance(content, StringTree) else StringTree('', content=content)
            # Make sure no existing grand child is updated. It would reside too
            # far up in the grand child OrderedDict, we need it last
            if gc.content in self.children[nibble].children:
                del self.children[nibble].children[gc.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() # type: ignore

    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, f'-- At "{self.content}": Adding "{path}" -> "{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, path = None, default=None) -> Optional[str]:
        if path:
            child = self.get(path)
            if child is None:
                if default:
                    return default
                return None
            return child.value()
        if len(self.children) == 0:
            raise Exception('tried to get value from leave "{}"'.format(self.content))
        slog(DEBUG, f'Returning value from children {self.children}')
        return self.children[next(reversed(self.children))].content # type: ignore

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

    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)