Convert Earley's chart to trees (brute force)

chart_to_trees.py

import json
import sys
import re
import pprint
import copy

pp = pprint.PrettyPrinter(indent=4)

class Node(object):
    def __init__(self, label, s=None, e=None, is_terminal=False, maxe=None):
        self.label = label
        self.s = s
        self.e = e
        self.children = []
        self.p = 0
        self.state = None
        self.is_terminal = is_terminal
        self.maxe = maxe

    def set_state(self, state):
        if self.s is not None and state["i"] != self.s:
            return False
        if self.e is not None and state["e"] != self.e:
            return False
        self.state = state
        self.e = state["e"]
        return True

    def update_children(self, tr):
        if self.state["type"] == "N":
            for i, sym in enumerate(self.state["r"]):
                child = Node(sym, 
                             self.state["i"] if i == 0 else None, 
                             self.state["e"] if i + 1 == len(self.state["r"]) else None)
                self.children.append(child)
                tr.add_q(child)
        else:
            sym = self.state["r"][0]
            child = Node(sym, self.state["i"], self.state["e"], True)
            self.children.append(child)

    def num_info(self, x):
        return str(x) if x is not None else "X"

    def info(self):
        return self.label+":"+self.num_info(self.s)+"_"+self.num_info(self.e)

class Tree(object):
    def __init__(self, root):
        self.root = root
        self.q = [root]

    def add_q(self, n):
        self.q.append(n)

    def has_q(self):
        return len(self.q) > 0

    def rm_q(self):
        self.q = self.q[1:]

    def jump_q(self, n):
        self.q = [n] + self.q

    def node(self):
        return self.q[0]

    def clone(self):
        return copy.deepcopy(self)
  

    def _to_str(self, n):
        if len(n.children) == 0:
            return n.info()
        else:
            children = " ".join(map(self._to_str, n.children))
            return ("(" + n.info() + " " + children + ")")

    def to_str(self):
        return self._to_str(self.root)

class TreesFinder(object):
    def __init__(self, chart):
        self.chart = list(
            map(lambda states: list(filter(lambda s: s["p"] == len(s["r"]), 
                                           states)), 
                chart))        

        self.wlen = len(self.chart) - 1
        self.root_state = [state for state in chart[self.wlen] 
                           if state["l"] == "ROOT"][0]
        self.make_i_index()

    def make_i_index(self):
        self.idx = [[] for i in range(self.wlen)]
        for e, states in enumerate(self.chart):
            for state in states:
                state["e"] = e
                self.idx[state["i"]].append(state)

    def find(self):
        root = Node("ROOT", 0, self.wlen)
        tr = Tree(root)
        self._find(tr)

    def get_states(self, n):
        def cond(state):
           if state["l"] != n.label:
                return False
           if n.e is not None and state["e"] != n.e:
                return False
           if n.maxe is not None and state["e"] > n.maxe:
                return False
           return True
        return filter(cond, self.idx[n.s])

    def _find(self, tr):
        #, [n.info() for n in tr.q])
        if not tr.has_q():
            print(tr.to_str())
            return
        n = tr.node()
        if n.state is None:
            if n.s is not None and n.e is not None and not n.s < n.e:
                return
            for state in self.get_states(n):
                _tr = tr.clone()
                _n = _tr.node()
                if _n.set_state(state):
                    _n.update_children(_tr)
                    self._find(_tr)
        else:
            if n.p < len(n.children):
                child = n.children[n.p]
                if n.p + 1 != len(n.children):
                    child.maxe = n.e - 1
                if n.p > 0:
                    child.s = n.children[n.p-1].e
                if child.is_terminal:
                    n.p += 1
                    self._find(tr.clone())
                else:
                    n.p += 1
                    tr.jump_q(child)
                    self._find(tr.clone())
            else:
                tr.rm_q()
                self._find(tr.clone())
def find_trees(chart):
    finder = TreesFinder(chart)
    finder.find()

for line in sys.stdin:
    chart = json.loads(re.split('[\r\n]', line)[0])
    find_trees(chart)