graph.py

from collections import defaultdict


class Graph:
    def __init__(self):
        self.nodes = defaultdict(set)

    def add_node(self, node, neighbors):
        self.nodes[node].update(neighbors)

        for neighbor in neighbors:
            self.nodes[neighbor].add(node)

    def __iter__(self):
        return iter(self.nodes)

    def traverse_graph(self, start_node, hook_function):
        visited = set()

        def dfs(node):
            if node not in visited:
                visited.add(node)
                hook_function(node)

                for neighbor in self.nodes[node]:
                    dfs(neighbor)

        dfs(start_node)

    def __str__(self):
        pretty_graph = ""
        for node, neighbors in self.nodes.items():
            pretty_graph += f"{node} -> {', '.join(str(n) for n in neighbors)}\n"
        return pretty_graph

type_inference.md

copy to dir type_inference

run python -m type_inference

type_inference.py

import typing
import graph


class Node:
    def __init__(
        self, name: str, kind: str, typeset: typing.Optional[list] = None
    ) -> None:
        self.name = name
        self.kind = kind

        self.typeset = None
        if typeset:
            self.typeset = set(typeset)

    def __repr__(self) -> str:
        typeset = self.typeset or "any"
        return f"Node({self.kind} {self.name} {typeset})"

    def __str__(self) -> str:
        return self.name


def infer_types(gr: graph.Graph):
    typed_nodes = [node for node in gr if node.typeset]

    for node in typed_nodes:

        def adjust_type(other):
            if other is node:
                return
            other.typeset = (other.typeset or node.typeset.copy()) & node.typeset

        gr.traverse_graph(node, adjust_type)


# x = 1
# y = x + 2
# z = sin(y)

const_1 = Node("1", "const", ["int", "float"])
const_2 = Node("2", "const", ["int", "float"])
x = Node("x", "var")
fn_plus = Node("+", "op", ["int", "float"])
y = Node("y", "var")
fn_sin = Node("sin", "op", ["float"])
z = Node("z", "var")


gr = graph.Graph()

gr.add_node(x, [fn_plus, const_1])
gr.add_node(fn_plus, [x, const_2, y])
gr.add_node(z, [fn_sin, y])

infer_types(gr)

for node in gr:
    print(f"{node.name} {node.typeset}")