graph shit.py

from functools import singledispatch, singledispatchmethod
import types, gc, re, sys
from IPython.core.formatters import (
    DisplayFormatter,
    BaseFormatter,
    catch_format_error,
    JSONFormatter,
)
from traitlets import ObjectName, Unicode, Instance, List, Any
from IPython import get_ipython
from pathlib import Path

TYPE, ID, GRAPH, CONTAINER, NEST, CONTEXT = "@type @id @graph @container @nest @context".split()
LD = "application/ld+json"
MAIN = __name__ == "__main__"
ACTIVE = "__file__" not in locals() and MAIN
shell = get_ipython()


class MetadataFormatter(BaseFormatter):
    (
        graph,
        format_type,
    ) = List(), Unicode("application/ld+json")
    _return_type, print_method = (list, dict), ObjectName("_repr_metadata_")

    @singledispatchmethod
    def get_id(self, object):
        return next(self.get_object(object), None)

    @singledispatchmethod
    def get_graph(self, object):
        data = {TYPE: self.get_id(type(object))}
        id = self.get_id(object)
        if id:
            data.setdefault(ID, id)
            if isinstance(id, list):
                return [{ID: x, **data} for x in id]
        return data

    def get_object(self, object, filter=None):
        if isinstance(filter, str):
            filter = re.compile(filter)
        for referrer in (x for x in gc.get_referrers(object) if isinstance(x, dict)):
            yield from self.get_object_from_ns(referrer, object, filter=filter)

    def get_object_from_ns(self, ns, object, filter=None):
        weakref = ns.get("__weakref__")
        parent = None
        if weakref:
            parent = self.get_id(weakref.__objclass__)
        else:
            parent = ns.get("__module__", ns.get("__name__"))
            if parent:
                parent += ":"

        if not parent and ns is sys.modules:
            return object.__name__

        for k in (k for k, v in ns.items() if v is object and not k.startswith("_")):
            name = f"{parent or 'noparent'}#{k}"
            if filter is not None and not filter.match(name):
                continue
            yield name

    def get_session_cell_id(self):
        data = get_ipython().kernel.get_parent()
        return data["metadata"]["cellId"], data["header"]["session"]

    def set_metadata(self, object=None, **kwargs):
        ids = dict(zip(("cell:id", "session:id"), self.get_session_cell_id()))
        ids.pop("session:id")
        node = {} if object is None else self.get_graph(object)
        if isinstance(node, dict):
            node = [node]
        for node in node:
            self.graph.append({**ids, **node, **kwargs})

    def __call__(self, object):
        explicit = super().__call__(object)
        if explicit:
            if isinstance(explicit, dict):
                self.set_metadata(**explicit)
            else:
                for e in explicit:
                    self.set_metadata(**w)
        else:
            self.set_metadata(object)
        try:
            return self.graph[:]
        finally:
            self.graph.clear()


class LinkedDataFormatter(DisplayFormatter):
    metadata_formatter = Instance(MetadataFormatter, args=())

    def format(self, object, include=None, exclude=None):
        data, meta = super().format(object, include, exclude)
        g = self.metadata_formatter(object)
        if g:
            meta[LD] = {GRAPH: g}
        return data, meta


def load_ipython_extension(shell=get_ipython()):
    shell.display_formatter = LinkedDataFormatter(**shell.display_formatter._trait_values)
    shell.user_ns["set_metadata"] = shell.display_formatter.metadata_formatter.set_metadata


def unload_ipython_extension(shell=get_ipython()):
    shell.display_formatter = DisplayFormatter(**shell.display_formatter._trait_values)


@MetadataFormatter.get_graph.register(tuple)
def get_graph_tuple(self, object):
    return list(map(self.get_graph, object))


@MetadataFormatter.get_id.register(types.ModuleType)
def get_name(self, object):
    return object.__name__


@MetadataFormatter.get_id.register(str)
def get_graph_str(self, object):
    from urllib.parse import urlparse

    parsed = urlparse(object)
    if parsed.scheme:
        return object