snake string

snake_string.py

"""S is a snake string

The goal of snake string to easily generate paramaterized strings from
attributes, args, and kwargs.

Examples:

S.asdf                   -> asdf
S.hey.how.are.you        -> hey.how.are.you
S.wow.s("this").is.cool  -> wow.this.is.cool
S.a.b.s(1,2,bluh=3).yay  -> a.b.1.2.3.yay
"""

import ast
import inspect
from functools import wraps
from collections import namedtuple

attr_size = namedtuple("attr_size", ["attr", "size"])

class _S_AttrVisitor(ast.NodeVisitor):
    """ast visitor to parse all attribute calls from a givin _S expression node
    """
    hist = []

    def clear(self):
        _S_AttrVisitor.hist = []

    def generic_visit(self, node):
        """Parse the "__getattr__" and "s" calls from the ast

        "__getattr__" calls are stored in Attribute nodes
        "s" calls are stored in Call nodes

        - End traversal if "_S" Name node is found
        - Add Attribute node names if they are not an "s" call
        - Add "s" calls only if they have args and kwargs
        """
        if isinstance(node, ast.Name):
            if globals().get(node.id, False).__class__.__name__ == _S.__name__:
                return
            _S_AttrVisitor.hist = []
        elif isinstance(node, ast.Attribute):
            if node.attr != _S.s.__name__:
                _S_AttrVisitor.hist.append(attr_size(node.attr, 1))
        elif isinstance(node, ast.Call):
            if isinstance(node.func, ast.Attribute):
                if node.func.attr == _S.s.__name__:
                    _S_AttrVisitor.hist.append(
                        attr_size(node.func.attr, len(node.args) + len(node.keywords)))
        
        ast.NodeVisitor.generic_visit(self, node)

class _S(type):
    """SnakeString metaclass

    When used as a metaclass the inheriting class should be interfaced via the
    __getattr__() or s() methods for making strings.

    Users should not interface with any dunder methods on this class or underlying
    mechanisms to creates attribute strings will break!
    """
    __accum = ""
    __parsed_attrs = []

    def __parse_and_add_sep(f):
        """decorator to pre parse attributes and add seperators
        """

        @wraps(f)
        def parse_and_add_sep(cls, *args, **kwargs):
            """wrapper function to check for following conditions

            1. Parse caller stack frame only once
            2. Get size of current and future attributes
            3. Check whether to add intra attribute seperators
            """

            ## Check the caller (outer stack frame) and parse expression
            ## using _S_AttrVisitor for caller attributes
            if not cls.__parsed_attrs:
                outer_frame = inspect.getouterframes(inspect.currentframe(), 2)[1]
                s_parse = _S_AttrVisitor()
                s_parse.visit(ast.parse(outer_frame.code_context[-1]).body[0])
                cls.__parsed_attrs = list(s_parse.hist)
                s_parse.clear()

            ## Check current attr and post accumulate future attrs
            current_attr_has_size = next(iter(cls.__parsed_attrs[::-1]), 0).size
            res = f(cls, *args, **kwargs)
            future_attrs_have_size = sum(attr.size for attr in cls.__parsed_attrs)

            ## Add seperator if current has size and future attrs have size
            if current_attr_has_size and future_attrs_have_size:
                cls.__accum += cls.sep
            return res
        return parse_and_add_sep

    @__parse_and_add_sep
    def s(cls, *args, **kwargs):
        """Add args and kwargs values to output string
        """

        ## If attributes have been pre parsed and the latest
        ## attribute is this function accumulate to output str
        if len(cls.__parsed_attrs) and (cls.__parsed_attrs[-1].attr == cls.s.__func__.__name__):
            cls.__parsed_attrs.pop()
            if list(args) + list(kwargs.values()):
                cls.__accum += cls.sep.join(map(str, list(args) + list(kwargs.values())))

        ## If there are no more attributes left to parse
        ## return the final string
        if len(cls.__parsed_attrs) == 0:
            return repr(cls)
        return cls

    @__parse_and_add_sep
    def __getattr__(cls, attr):
        """Custom getattr to accumulate the attr to the output string
        """

        ## If attributes have been pre parsed and the latest parsed
        ## attribute is the passed attribute, accmulate to the output str
        if (len(cls.__parsed_attrs) > 0) and (cls.__parsed_attrs[-1][0] == attr):
            cls.__parsed_attrs.pop()
            cls.__accum += attr

        ## If there are no more attributes left to parse
        ## return the final string
        if len(cls.__parsed_attrs) == 0:
            return repr(cls)
        return cls

    def __repr__(cls):
        """Once the string has been read, clear it
        """
        tmp = cls.__accum
        cls.__accum = ""
        cls.__parsed_attrs = []
        return str(tmp)
    
    def __str__(cls):
        """wrapper for repr
        """
        return repr(cls)
    
class SnakeString(metaclass=_S):
    sep = "."


S = SnakeString
print(S.woa)
print(S.a)
print(S.b.qw)
print(S.c.de.s(1))
print(S.a.b.s('c').d)
print(S.s('a.b.c.d').e.f)
print(S.a.b.c.d.s().e.f)
print(S.e.f.g.s(a=3).h)
print(S.metrics.score.s(5).yay)
print(S.yay == "yay")