Source code for feynml.feynml

import os
import warnings
from dataclasses import dataclass, field
from importlib.metadata import version
from typing import List, Optional

import smpl_io.io as io

from feynml.feynmandiagram import FeynmanDiagram
from feynml.head import Head
from feynml.id import Identifiable
from feynml.meta import Meta
from feynml.sheet import SheetHandler
from feynml.xml import XML

feynml_version = version("feynml")




@dataclass
class FeynML(SheetHandler, XML):
    class Meta:
        name = "feynml"

    version: Optional[str] = field(
        default=feynml_version, metadata={"name": "version", "type": "Attribute"}
    )

    # post init to check version
    def __post_init__(self):
        if self.version < feynml_version:
            warnings.warn(
                f"FeynML version {self.version} is older than this parser {feynml_version}."
            )
        elif self.version > feynml_version:
            warnings.warn(
                f"FeynML version {self.version} is newer than this parser {feynml_version}."
            )

        self.head.metas.append(Meta("feynml", version("feynml")))
        for d in self.diagrams:
            d.parent_fml = self

    head: Optional[Head] = field(
        default_factory=lambda: Head(),
        metadata={"name": "head", "namespace": "", "type": "Element"},
    )

    diagrams: List[FeynmanDiagram] = field(
        default_factory=list,
        metadata={"name": "diagram", "type": "Element", "namespace": ""},
    )

    def is_isomorphic(self, other: "FeynML"):
        if len(self.diagrams) != len(other.diagrams):
            return False
        found = [False] * len(self.diagrams)
        for d in self.diagrams:
            for i, od in enumerate(other.diagrams):
                if d.is_isomorphic(od):
                    if found[i]:
                        # No duplicates
                        return False
                    found[i] = True
        return all(found)

    def get_root(self):
        return self.head.root

    def with_root(self, root):
        self.head.root = root
        return self

    def get_diagram(self, idd):
        for d in self.diagrams:
            if d.id == idd:
                return d
        return None



    def get_style(self, obj, xml=None):
        if xml is None:
            xml = self
        for d in self.diagrams:
            if isinstance(obj, Identifiable):
                if d.has_id(obj.id):
                    return d.get_style(obj, xml)
        return self.head.get_style(obj, xml)




    def get_sheets(self):
        return self.head.get_sheets()




    def get_sheet(self):
        return self.head.get_sheet()




    def with_sheet(self, sheet):
        self.head.with_sheet(sheet)
        return self


    def render(self, show=True, **kwargs):
        for d in self.diagrams:
            d.render(show=show, **kwargs)

    def _ipython_display_(self):
        self.render(show=True)
        return self



    @classmethod
    def from_xml_file(cls, file: str):
        """Load self from XML file."""
        # get parent directory of file
        r = cls.from_xml(io.read(file))
        # check if file exists
        if os.path.isfile(file):
            r.head.root = os.path.dirname(file) + "/"
        elif file.startswith("http"):
            # remove last file in url
            r.head.root = "/".join(file.split("/")[:-1]) + "/"
        return r