Source code for scenic.simulators.webots.world_parser

"""Parser for WBT files using ANTLR.

The ANTLR parser itself, consisting of the *WBTLexer.py*, *WBTParser.py*, and *WBTVisitor.py* files, is autogenerated from *WBT.g4*.
"""

import io
import sys
import time

import antlr4
from antlr4.error.ErrorListener import ErrorListener
import numpy as np

from scenic.simulators.webots.WBTLexer import WBTLexer
from scenic.simulators.webots.WBTParser import WBTParser
from scenic.simulators.webots.WBTVisitor import WBTVisitor

### Setup

## Classes representing various types of VRML nodes


class Node:
    """A generic VRML node"""

    def __init__(self, nodeType, attrs):
        self.nodeType = nodeType
        self.attrs = attrs


## Classes for parsing with ANTLR

# Error handling


class ErrorReporter(ErrorListener):
    """ANTLR listener for reporting parse errors"""

    def syntaxError(self, recognizer, offendingSymbol, line, column, msg, e):
        raise Exception(f"Parse error at {line}:{column}: {msg}")


# Walking the parse tree and extracting values


class Evaluator(WBTVisitor):
    """Constructs an object representing the given value from the parse tree"""

    def __init__(self, nodeClasses):
        super().__init__()
        # Dictionary indicating which types of nodes have specialized classes
        self.nodeClasses = nodeClasses

    def visitNode(self, node):
        nodeType = node.Identifier().getText()
        attrs = self.visitNodeBody(node.nodeBody())
        if nodeType in self.nodeClasses:
            nodeClass = self.nodeClasses[nodeType]
            nodeObj = nodeClass(attrs)
        else:
            nodeObj = Node(nodeType, attrs)
        return nodeObj

    def visitNodeBody(self, body):
        return {
            attr.Identifier().getText(): self.visitValue(attr.value())
            for attr in body.attribute()
        }

    def visitValue(self, val):
        return self.visit(val.getChild(0))

    def visitVector(self, vec):
        return np.array([float(n.getText()) for n in vec.Number()])

    def visitVectorWithNewlines(self, vec):
        return self.visitVector(vec)

    def visitString(self, string):
        return string.getText()[1:-1]  # strip off quotes

    def visitArray(self, array):
        return [
            self.visit(val)
            for val in array.children
            if not isinstance(val, antlr4.tree.Tree.TerminalNode)
        ]

    def visitBoolean(self, boolean):
        return boolean.getText() == "TRUE"


def parse(path):
    """Parse a world from a WBT file"""
    # Open file
    charStream = antlr4.FileStream(path)

    # Set up lexer and parser
    errorListener = ErrorReporter()
    lexer = WBTLexer(charStream)
    lexer.addErrorListener(errorListener)
    tokenStream = antlr4.CommonTokenStream(lexer)
    parser = WBTParser(tokenStream)
    parser.addErrorListener(errorListener)

    # Parse file
    return parser.world()


def findNodeTypesIn(types, world, nodeClasses={}):
    """Find all nodes of the given types in a world"""
    evaluator = Evaluator(nodeClasses)
    types = tuple(types)
    instances = {ty: [] for ty in types}
    for node in world.children:
        if isinstance(node, WBTParser.DefnContext):
            node = node.node()
        nodeType = node.Identifier().getText()
        if nodeType in instances:
            val = evaluator.visitNode(node)
            instances[nodeType].append(val)
    return tuple(instances[ty] for ty in types)