pyne/datatypes.py

from dataclasses import dataclass
from enum import auto, Enum
from typing import Any
import re


class Datatype(Enum):
    NONE = auto()
    NULL = auto()
    COMMENT = auto()
    STRING = auto()
    INTEGER = auto()
    FLOAT = auto()
    COMPLEX = auto()


@dataclass
class Data:
    type: Datatype
    value: Any


@dataclass
class Variable:
    name: str
    type: Datatype
    value: Data

    def __setattr__(self, name, value):
        if name == "value" and isinstance(value, Data):
            if self.type == value.type:
                object.__setattr__(self, name, value)
            else:
                raise TypeError(
                    "Type inequality; cannot assign value of "
                    f"type {value.type} to variable of type {self.type}"
                )
        else:
            object.__setattr__(self, name, value)


NONE_TYPE = Variable(None, Datatype.NONE, None)

STR_MATCH_BEGIN = re.compile(r"^\"")
STR_MATCH_END = re.compile(r"\"$")
VAR_ASSIGN_TOKEN = re.compile(r"[A-Za-z]+ ?= ?")

VAR_MATCH = {
    re.compile(r"null"): Datatype.NULL,
    re.compile(r"\/\/.+"): Datatype.COMMENT,
    re.compile(r"\".+\""): Datatype.STRING,
    re.compile(r"-?\d+\.?\d{0,}?\+\d+\.?\d{0,}?i"): Datatype.COMPLEX,
    re.compile(r"-?\d+\.\d{0,}"): Datatype.FLOAT,
    re.compile(r"-?\d+"): Datatype.INTEGER,
}