Testing shape instructions for lottie-spec

test_rendering_instructions.py

# pip install lottie cairosvg
import io
import re
import sys
import enum
import math
import json
import lottie
import pathlib
import argparse


class RenderContext:
    def __init__(self):
        self.values = {}
        self.shape = lottie.objects.bezier.Bezier()
        self.shape.closed = True
        self.inputs = set()
        self.members = {}
        self.define_value("E_t", 0.5519150244935105707435627)

    def define_value(self, name, value):
        self.values[name] = value

    def add_vertex(self, vertex: lottie.NVector):
        self.shape.add_point(vertex)


class AstNode:
    def evaluate(self, context: RenderContext):
        raise NotImplementedError

    def get_inputs(self, context: RenderContext):
        return


class AstComment(AstNode):
    def __init__(self, comment):
        self.comment = comment

    def evaluate(self, context: RenderContext):
        pass


class AstBinaryNode(AstNode):
    def __init__(self, lhs: AstNode, rhs: AstNode):
        self.lhs = lhs
        self.rhs = rhs

    def get_inputs(self, context: RenderContext):
        self.lhs.get_inputs(context)
        self.rhs.get_inputs(context)


class AstUnaryNode(AstNode):
    def __init__(self, child: AstNode):
        self.child = child

    def get_inputs(self, context: RenderContext):
        self.child.get_inputs(context)


class AstBinOp(AstBinaryNode):
    def __init__(self, lhs: AstNode, rhs: AstNode, op: str):
        super().__init__(lhs, rhs)
        match op:
            case "cdot" | "times":
                op = "*"
            case "=":
                op = "=="
        self.op = op

    def evaluate(self, context: RenderContext):
        lhs = self.lhs.evaluate(context)
        rhs = self.rhs.evaluate(context)

        match self.op:
            case "+":
                return lhs + rhs
            case "-":
                return lhs - rhs
            case "*":
                return lhs * rhs
            case "/":
                return lhs / rhs
            case ">":
                return lhs > rhs
            case "<":
                return lhs < rhs
            case "==":
                return lhs == rhs

        raise Exception("Unknown operator %s" % self.op)


class AstUnOp(AstUnaryNode):
    def __init__(self, child: AstNode, op: str):
        super().__init__(child)
        self.op = op

    def evaluate(self, context: RenderContext):
        child = self.child.evaluate(context)

        match self.op:
            case "+":
                return child
            case "-":
                return -child
            case "floor":
                return math.floor(child)
            case "ceil":
                return math.ceil(child)

        raise Exception("Unknown operator %s" % self.op)


class AstStep(AstUnaryNode):
    def __init__(self, child: AstNode, op: str):
        super().__init__(child)
        self.op = op

    def evaluate(self, context: RenderContext):
        child = self.child.evaluate(context)

        match self.op:
            case "bez_v":
                context.add_vertex(child)
                return None
            case "bez_i":
                context.shape.in_tangents[-1] = child
                return None
            case "bez_o":
                context.shape.out_tangents[-1] = child
                return None

        raise Exception("Unknown operator %s" % self.op)


class AstLiteral(AstNode):
    def __init__(self, value):
        self.value = value

    def evaluate(self, context: RenderContext):
        return self.value


class AstPoint(AstBinaryNode):
    def evaluate(self, context: RenderContext):
        lhs = self.lhs.evaluate(context)
        rhs = self.rhs.evaluate(context)
        return lottie.Point(lhs, rhs)


class AstNaryNode(AstNode):
    def __init__(self, operands):
        self.operands = operands

    def get_inputs(self, context: RenderContext):
        for o in self.operands:
            o.get_inputs(context)


class AstNaryOp(AstNaryNode):
    def __init__(self, operands, op):
        super().__init__(operands)
        self.op = op

    def evaluate(self, context: RenderContext):
        operands = [o.evaluate(context) for o in self.operands]

        match self.op:
            case "min":
                return min(operands)
            case "max":
                return max(operands)

        raise Exception("Unknown operator %s" % self.op)


class AstBlock(AstNaryNode):
    def __init__(self, type):
        super().__init__([])
        self.type = type

    def evaluate(self, context: RenderContext):
        for o in self.operands:
            o.evaluate(context)


class AstNamedValue(AstNode):
    def __init__(self, name: str):
        self.name = name
        if "." in self.name:
            self.object, self.member = self.name.split(".")
        else:
            self.object = name
            self.member = None

    def evaluate(self, context: RenderContext):
        if self.member and self.name not in context.values:
            val = context.values[self.object]
            return getattr(val, self.member)
        return context.values[self.name]

    def get_inputs(self, context: RenderContext):
        if self.object not in context.values:
            context.inputs.add(self.name)
        if self.member:
            if self.object not in context.members:
                context.members[self.object] = {self.member}
            else:
                context.members[self.object].add(self.member)


class AstDefinition(AstUnaryNode):
    def __init__(self, child: AstNode, name: str):
        super().__init__(child)
        self.name = name

    def evaluate(self, context: RenderContext):
        context.values[self.name] = self.child.evaluate(context)

    def get_inputs(self, context: RenderContext):
        context.define_value(self.name, None)
        self.child.get_inputs(context)


class AstCondition(AstBinaryNode):
    def evaluate(self, context):
        if self.lhs.evaluate(context):
            self.rhs.evaluate(context)


class TokenType(enum.Enum):
    Operator = enum.auto()
    Literal = enum.auto()
    Lparen = enum.auto()
    Rparen = enum.auto()
    Lbrace = enum.auto()
    Rbrace = enum.auto()
    Name = enum.auto()
    Ignored = enum.auto()
    Eof = enum.auto()
    SlashName = enum.auto()
    Comma = enum.auto()
    Step = enum.auto()
    MathInline = enum.auto()
    MathBlock = enum.auto()


class Token:
    def __init__(self, type: TokenType, value=None):
        self.type = type
        self.value = value


class InstructionParser:
    def __init__(self, stream):
        self.stream = stream
        self.lookahead = Token(TokenType.Eof)

    def lex_getch_nospace(self) -> str:
        while True:
            c = self.stream.read(1)
            if not c.isspace():
                return c

    def lex_getch(self) -> str:
        return self.stream.read(1)

    def lex_unget(self):
        self.stream.seek(self.stream.tell() - 1)

    def lex_token(self) -> Token:
        while True:
            c = self.lex_getch_nospace()
            if c == '' or c not in "&":
                break

        if c == '':
            return Token(TokenType.Eof)

        if c == "$":
            d = self.lex_getch()
            if d == "$":
                return Token(TokenType.MathBlock)
            self.lex_unget()
            return Token(TokenType.MathInline)

        if c == '\\':
            d = self.lex_getch()
            if d == '\\':
                return self.lex_token()
            if d.isalpha():
                name = ""
                while d.isalpha():
                    name += d
                    d = self.lex_getch()
                self.lex_unget()

                if name in ("left", "right"):
                    return self.lex_token()

                if name in ("begin", "end"):
                    while True:
                        c = self.lex_getch()
                        if c == '}':
                            return self.lex_token()

                return Token(TokenType.SlashName, name)

        if c in "*+-=<>":
            return Token(TokenType.Operator, c)

        if c.isalpha():
            name = ""
            while c.isalpha() or c in "_.":
                name += c
                c = self.lex_getch()
            self.lex_unget()
            return Token(TokenType.Name, name)

        if c == "{":
            return Token(TokenType.Lbrace)

        if c == "}":
            return Token(TokenType.Rbrace)

        if c == "(":
            return Token(TokenType.Lparen)

        if c == ")":
            return Token(TokenType.Rparen)

        if c == ",":
            return Token(TokenType.Comma)

        # Skip markdown images
        if c == "!":
            d = self.lex_getch()
            if d == "[":
                while c != ")":
                    c = self.lex_getch()
                return self.lex()
            self.lex_unget()

        if c.isdigit():
            val = ""
            while c.isdigit() or c == ".":
                val += c
                c = self.lex_getch()

            self.lex_unget()
            if val.endswith("."):
                return Token(TokenType.Step)

            return Token(TokenType.Literal, float(val))

        raise Exception("Unknown token %s" % c)

    def lex(self):
        self.lookahead = self.lex_token()
        # print(self.stream.tell(), self.lookahead.type, self.lookahead.value)
        return self.lookahead

    def parse(self):
        try:
            self.lex()
            block = AstBlock("top")
            while self.lookahead.type != TokenType.Eof:
                block.operands.append(self.parse_statement(False))
            return block
        except Exception:
            pos = self.stream.tell()
            endl = pos
            while self.stream.read(1) not in ("", "\n"):
                endl += 1

            self.stream.seek(0)
            before = self.stream.read(endl)
            startl = before.rfind("\n") + 1

            col = pos - startl
            row = before.count("\n")
            sys.stderr.write("Error at position %s. Line %s Col %s\n" % (pos, row, col))
            sys.stderr.write(before[startl:endl])
            sys.stderr.write("\n%s^\n" % (" " * (pos-startl)))
            raise

    def parse_statement(self, in_condition):
        if self.lookahead.type == TokenType.Step:
            block = AstBlock("steps")
            while self.lookahead.type == TokenType.Step:
                block.operands.append(self.parse_step())
            return block if block.operands != 1 else block.operands[0]
        elif self.lookahead.type == TokenType.MathBlock:
            self.lex()
            block = AstBlock("math")
            while self.lookahead.type != TokenType.MathBlock:
                block.operands.append(self.parse_assignment())
            self.lex()
            return block if block.operands != 1 else block.operands[0]
        elif self.lookahead.type == TokenType.Name:
            comment = self.lookahead.value
            while True:
                c = self.lex_getch()
                if c == '' or c == '\n':
                    self.lex()
                    return AstComment(comment)

                if comment == "If":
                    pos = self.stream.tell()
                    if in_condition:
                        self.stream.seek(pos - 2)
                        return None
                    if self.lex().type == TokenType.MathInline:
                        self.lex()
                        return self.parse_condition_block()
                    else:
                        self.stream.seek(pos)

                comment += c
        else:
            self.expect(TokenType.Eof)

    def parse_condition_block(self):
        block = AstBlock("if")
        condition = AstCondition(self.parse_expression(), block)
        self.expect(TokenType.MathInline)
        comment = ""
        while True:
            c = self.lex_getch()
            if c == '' or c == '\n':
                self.lex()
                break
            comment += c
        comment = comment.strip(" .,")
        if comment:
            block.operands.append(AstComment(comment))

        while True:
            stmt = self.parse_statement(True)
            if not stmt:
                break
            block.operands.append(stmt)

        return condition

    def parse_step(self):
        instruction = ""
        c = self.lex_getch_nospace()
        while True:
            if c == '' or c == '\n':
                return AstComment(instruction)
            instruction += c
            match instruction:
                case "Add vertex":
                    self.lex()
                    return AstStep(self.parse_point(), "bez_v")
                case "Set in tangent":
                    self.lex()
                    return AstStep(self.parse_point(), "bez_i")
                case "Set out tangent":
                    self.lex()
                    return AstStep(self.parse_point(), "bez_o")

            c = self.lex_getch()

    def parse_assignment(self):
        name = self.expect(TokenType.Name).value
        self.expect(TokenType.Operator, "=")
        value = self.parse_expression()
        return AstDefinition(value, name)

    def parse_point(self):
        self.expect(TokenType.MathInline)
        self.expect(TokenType.Lparen)
        x = self.parse_expression()
        self.expect(TokenType.Comma)
        y = self.parse_expression()
        self.expect(TokenType.Rparen)
        self.expect(TokenType.MathInline)
        return AstPoint(x, y)

    def parse_expression(self):
        return self.parse_cmp()

    def has_operator(self, *op):
        return self.lookahead.type == TokenType.Operator and self.lookahead.value in op

    def has_slashname(self, *op):
        return self.lookahead.type == TokenType.SlashName and self.lookahead.value in op

    def parse_cmp(self):
        lhs = self.parse_add()
        while self.has_operator(">", "<", "="):
            op = self.lookahead.value
            self.lex()
            rhs = self.parse_add()
            lhs = AstBinOp(lhs, rhs, op)
        return lhs

    def parse_add(self):
        lhs = self.parse_mul()
        while self.has_operator("+", "-") or self.has_slashname("cdot", "times"):
            op = self.lookahead.value
            self.lex()
            rhs = self.parse_mul()
            lhs = AstBinOp(lhs, rhs, op)
        return lhs

    def parse_mul(self):
        lhs = self.parse_unary()
        while self.has_operator("*", "/"):
            op = self.lookahead.value
            self.lex()
            rhs = self.parse_unary()
            lhs = AstBinOp(lhs, rhs, op)
        return lhs

    def parse_unary(self):
        if self.lookahead.type == TokenType.Operator:
            op = self.lookahead.value
            self.lex()
            return AstUnOp(self.parse_expression(), op)

        if self.lookahead.type == TokenType.SlashName:
            name = self.lookahead.value
            if name in ("lfloor", "lceil"):
                self.lex()
                child = self.parse_expression()
                close = self.expect(TokenType.SlashName).value
                if close[1:] != name[1:]:
                    raise Exception("%s after %s" % (close, name))
                return AstUnOp(child, name[1:])
            elif name == "frac":
                self.lex()
                self.expect(TokenType.Lbrace)
                lhs = self.parse_expression()
                self.expect(TokenType.Rbrace)
                self.expect(TokenType.Lbrace)
                rhs = self.parse_expression()
                self.expect(TokenType.Rbrace)
                return AstBinOp(lhs, rhs, "/")
            elif name in ("min", "max"):
                self.lex()
                self.expect(TokenType.Lparen)
                operands = []
                while True:
                    operands.append(self.parse_expression())
                    if self.lookahead.type == TokenType.Rparen:
                        self.lex()
                        break
                    self.expect(TokenType.Comma)
                return AstNaryOp(operands, name)

        if self.lookahead.type == TokenType.Lparen:
            self.lex()
            child = self.parse_expression()
            self.expect(TokenType.Rparen)
            return child

        return self.parse_primary()

    def parse_primary(self):
        if self.lookahead.type == TokenType.Literal:
            node = AstLiteral(self.lookahead.value)
            self.lex()
            return node
        elif self.lookahead.type == TokenType.Name:
            node = AstNamedValue(self.lookahead.value)
            self.lex()
            return node
        self.raise_token()

    def expect(self, type, value=None):
        if self.lookahead.type != type:
            self.raise_token(type)

        if value is not None:
            if self.lookahead.value != value:
                raise Exception("Expected token value %r, got %r" % (value, self.lookahead.value))

        tok = self.lookahead
        self.lex()
        return tok

    def raise_token(self, expected=None):
        msg = "Unknown token %s %s" % (self.lookahead.type, self.lookahead.value)
        if expected:
            msg += ". Expected %s" % expected
        raise Exception(msg)


def ast_to_python(node: AstNode, indent=""):

    match node:
        case AstComment():
            return "%s# %s\n" % (indent, node.comment)
        case AstBinOp():
            return "(%s %s %s)" % (ast_to_python(node.lhs), node.op, ast_to_python(node.rhs))
        case AstUnOp():
            return "%s%s" % (node.op, ast_to_python(node.child))
        case AstLiteral():
            return repr(node.value)
        case AstPoint():
            return "NVector(%s, %s)" % (ast_to_python(node.lhs), ast_to_python(node.rhs))
        case AstNaryOp():
            return "%s(%s)" % (node.op, ", ".join(map(ast_to_python, node.operands)))
        case AstBlock():
            return "\n".join(ast_to_python(p, indent) for p in node.operands) + "\n"
        case AstNamedValue():
            return node.name
        case AstDefinition():
            return "%s%s = %s" % (indent, node.name, ast_to_python(node.child, indent))
        case AstStep(op="bez_v"):
            return "%sshape.add_point(%s)" % (indent, ast_to_python(node.child))
        case AstStep(op="bez_i"):
            return "%sshape.in_tangents[-1] = %s" % (indent, ast_to_python(node.child))
        case AstStep(op="bez_o"):
            return "%sshape.out_tangents[-1] = %s" % (indent, ast_to_python(node.child))
        case AstCondition():
            next_indent = indent + (" " * 4)
            return "%sif %s:\n%s" % (indent, ast_to_python(node.lhs), ast_to_python(node.rhs, next_indent))
        case _:
            raise Exception("Unknown node %s" % node)


def ast_to_json(node):
    if isinstance(node, AstNode):
        data = {"_": type(node).__name__}
        for k, v in vars(node).items():
            data[k] = ast_to_json(v)
        return data
    elif isinstance(node, list):
        return [ast_to_json(v) for v in node]
    return node


class InputReader:
    def __init__(self, context: RenderContext, defaults: dict):
        self.context = context

        for key, val in defaults.items():
            if isinstance(val, list):
                val = lottie.NVector(*val)
            context.define_value(key, val)
            context.inputs.add(key)

    def get_value(self, var):
        hint = ""
        if self.context.members.get(var, set()) & {"x", "y"}:
            hint = " (point)"

        strval = input("%s%s: " % (var, hint))
        if "," in strval:
            return lottie.NVector(*map(float, strval.split(",")))
        else:
            return float(strval)

    def resolve(self):
        for var in sorted(context.inputs):
            if var not in context.values:
                context.define_value(var, self.get_value(var))


def render_lottie(context: RenderContext):
    ast.evaluate(context)
    an = lottie.objects.Animation()
    lay = lottie.objects.ShapeLayer()
    an.add_layer(lay)
    path = lottie.objects.Path(context.shape)
    bbox = path.bounding_box()
    margin = 10
    an.width = bbox.width + 2 * margin
    an.height = bbox.height + 2 * margin
    lay.shapes.append(path)
    lay.shapes.append(lottie.objects.Stroke(lottie.Color(1, 0, 0), 4))
    lay.transform.position.value = -lottie.NVector(bbox.x1 - margin, bbox.y1 - margin)
    return an


def render_python(context: RenderContext, fp):
    fp.write("import lottie\n")
    fp.write("from lottie import NVector\n")

    cleaned_inputs = set()
    for input in context.inputs:
        cleaned_inputs.add(input)
        if "." in input:
            cleaned_inputs.add(input.split(".")[0])

    input_str = "\n# Inputs\n"
    needs_mock = False
    for input in sorted(cleaned_inputs):
        if input in context.values:
            input_str += "%s = %r\n" % (input, context.values[input])
        else:
            input_str += "%s = unittest.mock.MagicMock()\n" % input
            needs_mock = True

    if needs_mock:
        fp.write("import unittest.mock\n")

    fp.write(input_str)

    fp.write("\nE_t = 0.5519150244935105707435627\n\n")
    fp.write("shape = lottie.objects.bezier.Bezier()\nshape.closed = True\n\n")

    fp.write(ast_to_python(ast))


class File:
    def __init__(self, file):
        self.file = sys.stdout if file is True else open(file, "w")

    def __enter__(self):
        return self.file

    def __exit__(self, *a):
        if self.file is not sys.stdout:
            self.file.close()


if __name__ == "__main__":
    arg_parser = argparse.ArgumentParser()
    file_arg = dict(nargs="?", type=pathlib.Path, const=True)
    arg_parser.add_argument("--shape", "-s", default="rectangle")
    arg_parser.add_argument("--json", "-j", **file_arg)
    arg_parser.add_argument("--python", "-py", **file_arg)
    arg_parser.add_argument("--lottie", "-l", **file_arg)
    arg_parser.add_argument("--png", type=pathlib.Path)
    arg_parser.add_argument("--defaults", "-d", type=str, default="")
    args = arg_parser.parse_args()

    shapes_md_path = pathlib.Path(__file__).absolute().parent.parent / "docs" / "specs" / "shapes.md"
    with open(shapes_md_path) as file:
        shapes_md = file.read()
    start = re.search("<h[0-9][^>]*id=['\"]%s['\"]" % args.shape, shapes_md).end()
    shapes_md = shapes_md[start:]
    start = re.search("lottie-playground>", shapes_md).end()
    shapes_md = shapes_md[start:]
    end = re.search("^<h[0-9]", shapes_md, re.MULTILINE).start()

    stream = io.StringIO(shapes_md[:end])
    defaults = {}
    if args.defaults.startswith("{"):
        defaults = json.loads(args.defaults)
    elif args.defaults:
        with open(args.defaults) as f:
            defaults = json.load(f)

    parser = InstructionParser(stream)
    ast = parser.parse()
    if args.json:
        with File(args.json) as f:
            f.write(json.dumps(ast_to_json(ast), indent=4) + "\n")

    context = RenderContext()
    inputs = InputReader(context, defaults)
    ast.get_inputs(context)

    if args.lottie:
        inputs.resolve()
        an = render_lottie(context)
        with File(args.lottie) as f:
            f.write(json.dumps(an.to_dict(), indent=4) + "\n")
    else:
        an = None

    if args.png:
        if an is None:
            inputs.resolve()
            an = render_lottie(context)
        with open(args.png, "wb") as fp:
            lottie.exporters.cairo.export_png(an, fp)

    if args.python:
        with File(args.python) as f:
            render_python(context, f)