Source code for fosf.parsers.similarity

#!/usr/bin/env python3

import networkx as nx
from lark import Lark

from fosf.config import FOSF_GRAMMAR
from fosf.parsers.base import BaseOSFParser, _BaseOSFTransformer
from fosf.syntax.similarity import Similarity


class _SimilarityTransformer(_BaseOSFTransformer):
    def __init__(self):
        super().__init__()
        self.sim = []

    def sim_equation(self, tree):
        "Handle a similarity equation."
        alpha = float(tree[-1].value)
        for i, _ in enumerate(tree[:-2]):
            sort_u = tree[i]
            sort_v = tree[i + 1]
            self.sim.append((sort_u, sort_v, alpha))

    def sim_equations(self, tree):
        g = nx.Graph()
        for u, v, w in self.sim:
            g.add_edge(u, v, weight=w)
        return Similarity(g, alpha=0, strict=True)

    def transform(self, tree):
        self.sim = []
        return super().transform(tree)




class SimilarityParser(BaseOSFParser):
    def __init__(self):
        self.parser = Lark.open_from_package(
            "fosf", FOSF_GRAMMAR, start="sim_equations"
        )
        self.transformer = _SimilarityTransformer()



    def parse(self, expression: str) -> Similarity:
        parse_tree = self.parser.parse(expression)
        return self.transformer.transform(parse_tree)