"""Generic text parsers for conversion of lattices in different formats to AT"""
from __future__ import annotations
__all__ = [
"AnyDescr",
"ElementDescr",
"SequenceDescr",
"BaseParser",
"UnorderedParser",
"DictNoDot",
]
from os import getcwd
from os.path import join, normpath, dirname
import re
import abc
from itertools import repeat
from collections.abc import Callable, Iterable, Generator, Mapping
from .utils import split_ignoring_parentheses, protect, restore
from ..lattice import Element, Lattice, params_filter
_dot = re.compile(r"\.?[a-z][\w.]*") # An identifier: starts with a letter
_singlequoted = re.compile(r"'([\w.]*)'")
_named = re.compile(r"name=([\w.]*)")
_doublequoted = re.compile(r'^"(.*)"$')
def _default_arg_parser(parser: BaseParser, argstr: str):
"""
Evaluate a keyword argument of a command and return the pair (key, value)
"""
try:
key, value = split_ignoring_parentheses(argstr, delimiter="=")
except ValueError: # Positional argument -> boolean flag
argstr = argstr.lower()
if argstr.startswith("-"):
v = False
key = argstr[1:]
else:
key = argstr
v = True
else: # Keyword argument
key = key.lower().replace("from", "frm")
# noinspection PyProtectedMember
if key in parser._str_arguments:
v = _doublequoted.sub(r"\1", value)
else:
v = parser.evaluate(value)
return key, v
[docs]
class DictNoDot(dict):
@staticmethod
def _no_dot(expr):
def repl(match):
return match.group().replace(".", "_")
return _dot.sub(repl, expr.lower())
def __setitem__(self, key, value):
super().__setitem__(self._no_dot(key), value)
def __getitem__(self, key):
return super().__getitem__(self._no_dot(key))
[docs]
def get(self, key, *args):
return super().get(self._no_dot(key), *args)
[docs]
class AnyDescr(abc.ABC):
"""Base class for source object descriptors"""
def __init__(self, *args, **kwargs):
self.name = kwargs.pop("name", self.__class__.__name__)
self.inverse = kwargs.pop("inverse", False)
# kwargs.setdefault("madclass", self.__class__.__name__)
super().__init__(*args, **kwargs)
def __neg__(self):
return self.inverted(copy=True)
def __call__(self, *args, copy: bool = True, **kwargs) -> AnyDescr | None:
"""Create a copy of the element with updated fields"""
if copy:
b = {key: kwargs.pop(key, value) for key, value in vars(self).items()}
b.update(kwargs)
# b.update(madclass=self.name)
return type(self)(self, *args, **b)
else:
self.update(*args, **kwargs)
return None
[docs]
def update(self, *args, **kwargs):
# Update attributes
for key, value in vars(self).items():
setattr(self, key, kwargs.pop(key, value))
if isinstance(self, Mapping):
# Update mapping
super().update(*args, **kwargs)
else:
# Add new attributes
for key, value in kwargs:
setattr(self, key, value)
[docs]
def inverted(self, copy=False):
"""Return a reversed element or line"""
instance = self() if copy else self
instance.inverse = not self.inverse
return instance
[docs]
@abc.abstractmethod
def expand(self, parser: BaseParser) -> Generator[Element, None, None]:
"""Iterator on the possibly inverted sequence of elements"""
pass
[docs]
class ElementDescr(AnyDescr, dict):
"""Simple representation of an element as a dict"""
def __init__(self, *args, **kwargs):
kwargs.pop("copy", False)
kwargs.setdefault("madtype", self.__class__.__name__)
super().__init__(*args, **kwargs)
def __getattr__(self, item):
# Allows accessing items using the attribute access syntax
return self[item]
def __rmul__(self, other):
"""Element repetition"""
return list(repeat(self, other))
def __repr__(self):
descr = super().copy()
cls = descr.pop("madtype")
keywords = [f"name={self.name}"]
keywords += [f"{k}={v!r}" for k, v in descr.items()]
return f"{cls}({', '.join(keywords)})"
[docs]
@staticmethod
def convert(name: str, *args, **params) -> list[Element]:
"""Generate the AT element, Most be overloaded for each specific element"""
return []
# noinspection PyUnusedLocal
[docs]
def expand(self, parser: BaseParser) -> Generator[Element, None, None]:
"""Iterator on the possibly inverted sequence of elements"""
try:
elems = self.convert(self.name, **self)
except Exception as exc:
exc.args += (f"{self}",)
raise
if self.inverse:
for elem in reversed(elems):
yield elem.swap_faces(copy=True)
else:
yield from elems
@property
def length(self) -> float:
"""Element length"""
return self.get("l", 0.0)
[docs]
class SequenceDescr(AnyDescr, list, abc.ABC):
"""Simple representation of a sequence of elements as a list"""
def __repr__(self):
str = super().__repr__()
return f"{self.__class__.__name__}({str})"
@property
def length(self) -> float:
return getattr(self, "l", 0.0)
[docs]
class BaseParser(DictNoDot):
"""Generic file parser
Analyses files with the following MAD-like format:
``variable = value``
``label : command [,attribute=value] [,attribute=value]...``
The parser builds a database of all the defined objects
"""
_str_arguments = set()
_argument_parser = {}
def __init__(
self,
env: dict,
*args,
delimiter: str | None = None,
continuation: str = "\\",
linecomment: str | tuple[str] | None = "#",
blockcomment: tuple[str, str] | None = None,
endfile: str | None = None,
verbose: bool = False,
**kwargs,
):
"""
Args:
env: global namespace used for evaluating commands
delimiter: command delimiter
continuation: command continuation character
linecomment: Line comment character
blockcomment: Block comment delimiter
endfile: "End of input" marker
verbose: If True, print detail on the processing
*args: dict initializer
**kwargs: dict initializer
"""
if isinstance(linecomment, tuple):
def line_comment(line):
for linecom in linecomment:
line, *_ = line.split(sep=linecom, maxsplit=1)
return line
else:
if linecomment is None:
def line_comment(line):
return line
else:
def line_comment(line):
line, *_ = line.split(sep=linecomment, maxsplit=1)
return line
if blockcomment is None:
# noinspection PyUnusedLocal
def handle_comments(buffer, line, in_comment):
buffer.append(line_comment(line))
return False, ""
else:
def handle_comments(buffer, line, in_comment):
if in_comment:
*rest, line = line.split(sep=endcomment, maxsplit=1)
in_comment = len(rest) <= 0
return in_comment, "" if in_comment > 0 else line
else:
line = line_comment(line)
contents, *rest = line.split(sep=begcomment, maxsplit=1)
buffer.append(contents)
in_comment = len(rest) > 0
return in_comment, rest[0] if in_comment else ""
begcomment, endcomment = blockcomment
self.skip_comments = handle_comments
self.delimiter = delimiter
self.continuation = continuation
self.endfile = endfile
self._verbose = verbose
self.env = env
self.bases = [getcwd()]
self.kwargs = kwargs
super().__init__(*args, **kwargs)
def _print(self, *args, **kwargs):
if self._verbose:
print(*args, **kwargs)
[docs]
def clear(self):
"""Clean the database"""
super().clear()
self.update(self.kwargs)
[docs]
def evaluate(self, expr):
"""Evaluate an expression using *self* as local namespace"""
return eval(expr, self.env, self)
def _eval_cmd(self, cmdname: str, no_global: bool = False) -> Callable:
"""Evaluate a command"""
cmd: Callable | None = self.get(cmdname, None)
if cmd is not None:
return cmd
else:
cmdname = cmdname.lower()
cmd = self.env.get(cmdname, None)
if cmd is None:
raise KeyError(cmdname)
elif no_global:
raise TypeError(f"{cmdname!r} is not allowed in this context")
else:
return cmd
@staticmethod
def _reason(exc: Exception) -> str:
"""Extract the element name from the exception"""
if isinstance(exc, KeyError): # Undefined element, attribute
return exc.args[0]
elif isinstance(exc, NameError): # refpos missing
return _singlequoted.search(exc.args[0])[1]
elif isinstance(exc, TypeError):
idx = _named.search(exc.args[-1]) # Missing pos. arg.
if idx is None:
idx = _singlequoted.search(exc.args[0]) # Not allowed in seq.
if idx is None:
return "TypeError"
else:
return idx[1]
elif isinstance(exc, ValueError): # overlap
print(exc.args[0])
return _singlequoted.search(exc.args[0])[1]
else:
return type(exc).__name__
def _argparser(self, command: str, *args: str):
argparser = self._argument_parser.get(command.lower(), _default_arg_parser)
return (argparser(self, arg) for arg in args)
def _assign(self, label: str, key: str, value: str):
"""Variable assignment"""
return key, self.evaluate(value)
def _raw_command(
self,
label: str | None,
cmdname: str,
*argnames: str,
no_global: bool = False,
**kwargs,
):
"""Command execution"""
func = self._eval_cmd(cmdname, no_global=no_global)
kwargs.update(self._argparser(cmdname, *argnames))
if label is None:
kwargs.setdefault("copy", False)
else:
kwargs.setdefault("name", label)
return func(**kwargs)
def _command(self, *args, **kwargs):
return self._raw_command(*args, **kwargs)
def _format_statement(self, line: str) -> str:
"""Reformat the input line
Overload this method for specific languages"""
line, matches = protect(line, fence=('"', '"')) # protect the quoted parts
line = "".join(line.split()) # Remove all spaces
(line,) = restore(matches, line)
return line
def _statement(self, line: str) -> bool:
"""Split a statement in 'label: command'"""
fmtline = self._format_statement(line)
if self.endfile is not None and fmtline.startswith(self.endfile):
return False
label, *cmd = fmtline.split(":", maxsplit=1)
if cmd: # label
fmtline = cmd[0]
else: # no label
label = None
arguments = split_ignoring_parentheses(fmtline)
self._decode(label, *arguments)
return True
def _decode(self, label: str, cmdname: str, *argnames: str) -> None:
"""Execute the split statement"""
left, *right = cmdname.split("=")
if right:
label, result = self._assign(label, left, right[0])
else:
result = self._command(label, cmdname, *argnames)
if not (label is None or result is None):
self[label] = result
def _finalise(self, final: bool = True) -> None:
"""Called at the end of processing"""
pass
@staticmethod
def _command_str(_, label, cmdname, *argnames):
string = ", ".join((cmdname, *argnames))
if label is not None:
string = " : ".join((label, string))
return string
def _analyse(self, key: str) -> None:
"""Print info on failed statements"""
print(f"\n{key!r} is not defined\n")
[docs]
def expand(self, key: str) -> Generator[Element, None, None]:
"""iterator over AT objects generated by a source object"""
try:
v = self[key]
if isinstance(v, AnyDescr):
yield from v.expand(self)
else:
yield v
except Exception as exc:
print(f"{type(exc).__name__}: {exc.args[0]}")
for arg in exc.args[1:]:
print(arg)
self._analyse(self._reason(exc))
raise
def _generator(self, params):
"""Generate AT elements for the Lattice constructor"""
use = params.setdefault("use", "ring")
params.setdefault("name", use)
# Iterate from the elements
yield from self.expand(use)
[docs]
def lattice(self, use="ring", **kwargs):
"""Create a lattice from the selected sequence
Parameters:
use: Name of the MADX sequence or line containing the desired
lattice. Default: ``ring``
Keyword Args:
name (str): Name of the lattice. Default: MADX sequence name.
particle(Particle): Circulating particle. Default: from MADX
energy (float): Energy of the lattice [eV], Default: from MADX
periodicity(int): Number of periods. Default: 1
*: All other keywords will be set as Lattice attributes
"""
return Lattice(self._generator, iterator=params_filter, use=use, **kwargs)
[docs]
def parse_lines(
self,
lines: Iterable[str],
final: bool = True,
**kwargs,
) -> None:
"""Process input lines and fill the database
Args:
lines: Iterable of input lines
final: If :py:obj:`True`, signals that the undefined variables may be set
to the default value
**kwargs: Initial variable definitions
"""
self.update(**kwargs)
buffer = []
in_comment: bool = False
ok: bool = True
for line_number, contents in enumerate(lines):
# Handle comments
while contents:
in_comment, contents = self.skip_comments(buffer, contents, in_comment)
if buffer:
contents = "".join(buffer).strip()
buffer = []
if not contents:
continue
else:
continue
# print(repr(contents))
# Handle delimiters
if self.delimiter is None:
statements = []
last = contents
else:
*statements, last = contents.split(sep=self.delimiter)
# Handle continuation
if self.continuation is None:
buffer.append(last)
else:
idc = last.find(self.continuation)
if idc >= 0:
buffer.append(last[:idc])
else:
statements.append(last)
# Process statements
for stmnt in statements:
stmnt = stmnt.strip()
if not stmnt:
continue
try:
ok = self._statement(stmnt)
except Exception as exc:
message = f"Line {line_number} {stmnt!r}, {exc}"
raise type(exc)(message) from exc
if not ok:
break
if not ok:
break
self._finalise(final=final)
[docs]
def parse_files(
self,
*filenames: str,
final: bool = True,
prolog: None | int | Callable[..., None] = None,
epilog: Callable[..., None] | None = None,
**kwargs,
) -> None:
"""Process files and fill the database
Args:
*filenames: Files to process
final: If :py:obj:`True`, signals that the undefined variables may be set
to the default value
prolog:
epilog:
**kwargs: Initial variable definitions
"""
self.update(**kwargs)
last = len(filenames) - 1
for nf, fn in enumerate(filenames):
fn = normpath(join(self.bases[-1], fn))
self.bases.append(dirname(fn))
print("Processing", fn)
try:
with open(fn) as f:
if callable(prolog):
prolog(f)
elif isinstance(prolog, int):
for _i in range(prolog):
next(f)
self.parse_lines(f, final=final and (nf == last))
if callable(epilog):
epilog(f)
finally:
print("End", fn)
self.bases.pop()
[docs]
class UnorderedParser(BaseParser):
"""parser allowing definitions in any order
This is done by storing the failed statements in a queue and iteratively trying
to execute them after all input statements have been processed, until the number
of failures is constant (hopefully zero)
"""
def __init__(self, env: dict, *args, **kwargs):
"""
Args:
env: global namespace
delimiter: command delimiter
continuation: command continuation character
linecomment: Line comment character
blockcomment: Block comment delimiter
endfile: End of input marker
verbose: If True, print detail on the processing
*args: dict initializer
**kwargs: dict initializer
"""
super().__init__(env, *args, **kwargs)
self.delayed = []
[docs]
def clear(self):
super().clear()
self.delayed = []
def _postpone(self, reason, label, cmdname, *argnames):
"""Store failing commands in self.delayed for later evaluation"""
self.delayed.append((reason, label, cmdname, *argnames))
def _decode(self, label: str, cmdname: str, *argnames: str) -> None:
"""Postpone failing commands"""
try:
super()._decode(label, cmdname, *argnames)
except (KeyError, NameError) as exc: # store the failing assignment
self._postpone(self._reason(exc), label, cmdname, *argnames)
def _finalise(self, final: bool = True) -> None:
"""Loop on evaluation of the pending statements"""
nend = len(self.delayed)
if nend > 0:
self._print(f"\nDelayed evaluation of {nend} statements\n")
while nend > 0:
statements = self.delayed
self.delayed = []
nstart = nend
for _reason, *args in statements:
self._decode(*args)
nend = len(self.delayed)
if nend == nstart:
break
def _lookup(self, item: str):
"""Search for an object in the pending statements"""
for reason, label, *args in self.delayed:
if label is not None and label.lower() == item:
return (reason, label, *args)
return None
def _missing(self, verbose: bool = False):
miss = set()
for cmd in self.delayed:
reason = cmd[0]
if reason == cmd[2].lower():
if verbose:
print(
f"Unknown command {cmd[2]!r} ignored: "
f"{self._command_str(*cmd)!r}"
)
continue
while cmd is not None:
reason = cmd[0]
cmd = self._lookup(reason)
miss.add(reason)
return miss
def _analyse(self, key: str) -> None:
"""Print the chain of failing commands"""
cmd = self._lookup(key.lower())
reason = key
while cmd is not None:
reason = cmd[0]
print(f"\n{key!r} depends on {reason!r}: {self._command_str(*cmd)!r}")
key = reason
cmd = self._lookup(reason)
print(f"\n{reason!r} is not defined\n")
missing = property(_missing, doc="Set of missing definitions")
