3676917429

Build Type

push

github-actions

Committed by Thomas Gläßle

Commit Message

Add documentation for building on linux+conda

Run Details

1045 of 1117 relevant lines covered (93.55%)

9.28 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

96.62

/src/cpymad/util.py

"""
Utility functions used in other parts of the pymad package.
"""
import re
import os
import sys
import tempfile
from collections import namedtuple
from contextlib import contextmanager
from enum import Enum
from numbers import Number
import collections.abc as abc

import numpy as np

from cpymad.parsing import Parser
from cpymad.types import (
    Range, Constraint,
    PARAM_TYPE_LOGICAL, PARAM_TYPE_INTEGER,
    PARAM_TYPE_DOUBLE, PARAM_TYPE_STRING, PARAM_TYPE_CONSTRAINT,
    PARAM_TYPE_LOGICAL_ARRAY, PARAM_TYPE_INTEGER_ARRAY,
    PARAM_TYPE_DOUBLE_ARRAY, PARAM_TYPE_STRING_ARRAY)


__all__ = [
    'mad_quote',
    'is_identifier',
    'name_from_internal',
    'name_to_internal',
    'format_param',
    'format_cmdpar',
    'format_command',
    'check_expression',
    'temp_filename',
    'ChangeDirectory',
]


# In CPython 3.6 dicts preserve insertion order (until deleting an element)
# Although, this is considered an implementation detail that should not be
# relied upon, we do so anyway:
ordered_keys = dict.keys if sys.version_info >= (3, 6) else sorted


def mad_quote(value: str) -> str:
    """Add quotes to a string value."""
    if '"' not in value:
        return '"' + value + '"'
    if "'" not in value:
        return "'" + value + "'"
    # MAD-X doesn't do any unescaping (otherwise I'd simply use `json.dumps`):
    raise ValueError("MAD-X unable to parse string with escaped quotes: {!r}"
                     .format(value))


def _fix_name(name: str) -> str:
    if name.startswith('_'):
        raise AttributeError("Unknown item: {!r}! Did you mean {!r}?"
                             .format(name, name.strip('_') + '_'))
    if name.endswith('_'):
        name = name[:-1]
    return name


# precompile regexes for performance:
re_compile = lambda s: re.compile(s, re.IGNORECASE)
_re_is_identifier = re_compile(r'^[a-z_][a-z0-9_]*$')
_re_symbol = re_compile(r'([a-z_][a-z0-9._]*(->[a-z_][a-z0-9._]*(\[[0-9]+\])?)?)')
_re_element_internal = re_compile(r'^([a-z_][a-z0-9_.$]*)(:\d+)?$')
_re_element_external = re_compile(r'^([a-z_][a-z0-9_.$]*)(\[\d+\])?$')


def is_identifier(name: str) -> bool:
    """Check if ``name`` is a valid identifier in MAD-X."""
    return bool(_re_is_identifier.match(name))


def expr_symbols(expr: str) -> set:
    """
    Return all symbols names used in an expression.

    For now this includes not only variables but also element attributes (e.g.
    ``quad->k1``) as well as function names (e.g. ``sin``).
    """
    return {m[0] for m in _re_symbol.findall(expr)}


def name_from_internal(element_name: str) -> str:
    """
    Convert element name from internal representation to user API. Example:

    >>> name_from_internal("foo:1")
    foo
    >>> name_from_internal("foo:2")
    foo[2]

    Element names are stored with a ":d" suffix by MAD-X internally (data in
    node/sequence structs), but users must use the syntax "elem[d]" to access
    the corresponding elements. This function is used to transform any string
    coming from the user before passing it to MAD-X.
    """
    try:
        name, count = _re_element_internal.match(element_name).groups()
    except AttributeError:
        raise ValueError("Not a valid MAD-X element name: {!r}"
                         .format(element_name)) from None
    if count is None or count == ':1':
        return name
    return name + '[' + count[1:] + ']'


def _parse_element_name(element_name: str) -> tuple:
    """
    Parse element name from user API. Example:

    >>> _parse_element_name("foo")
    (foo, None)
    >>> _parse_element_name("foo[2]")
    (foo, 2)

    See :func:`name_from_internal' for further information.
    """
    try:
        name, count = _re_element_external.match(element_name).groups()
    except AttributeError:
        raise ValueError("Not a valid MAD-X element name: {!r}"
                         .format(element_name)) from None
    if count is None:
        return name, None
    return name, int(count[1:-1])


def name_to_internal(element_name: str) -> str:
    """
    Convert element name from user API to internal representation. Example:

    >>> name_to_external("foo")
    foo:1
    >>> name_to_external("foo[2]")
    foo:2

    See :func:`name_from_internal` for further information.
    """
    name, count = _parse_element_name(element_name)
    return name + ':' + str(1 if count is None else count)


def normalize_range_name(name: str, elems=None) -> str:
    """Make element name usable as argument to the RANGE attribute."""
    if isinstance(name, tuple):
        return tuple(map(normalize_range_name, name))
    if '/' in name:
        return '/'.join(map(normalize_range_name, name.split('/')))
    name = name.lower()
    if name.endswith('$end') or name.endswith('$start'):
        if elems is None:
            return u'#s' if name.endswith('$start') else u'#e'
        else:
            return u'#s' if elems.index(name) == 0 else u'#e'
    return name


QUOTED_PARAMS = {'file', 'halofile', 'sectorfile', 'trueprofile'
                 'pipefile', 'trackfile', 'summary_file', 'filename',
                 'echo', 'title', 'text', 'format', 'dir'}


def format_param(key: str, value) -> str:
    """
    Format a single MAD-X command parameter.

    This is the old version that does not use type information from MAD-X. It
    is therefore not limited to existing MAD-X commands and attributes, but
    also less reliable for producing valid MAD-X statements.
    """
    if value is None:
        return None
    key = _fix_name(str(key).lower())
    if isinstance(value, Constraint):
        constr = []
        if value.min is not None:
            constr.append(key + '>' + str(value.min))
        if value.max is not None:
            constr.append(key + '<' + str(value.max))
        if constr:
            return u', '.join(constr)
        else:
            return key + '=' + str(value.val)
    elif isinstance(value, bool):
        return key + '=' + str(value).lower()
    elif key == 'range' or isinstance(value, Range):
        return key + '=' + _format_range(value)
    # check for basestrings before abc.Sequence, because every
    # string is also a Sequence:
    elif isinstance(value, str):
        if key in QUOTED_PARAMS:
            return key + '=' + mad_quote(value)
        else:
            # MAD-X parses strings incorrectly, if followed by a boolean.
            # E.g.: "beam, sequence=s1, -radiate;" does NOT work! Therefore,
            # these values need to be quoted. (NOTE: MAD-X uses lower-case
            # internally and the quotes prevent automatic case conversion)
            return key + '=' + mad_quote(value.lower())
    # don't quote expressions:
    elif isinstance(value, str):
        return key + ':=' + value
    elif isinstance(value, abc.Sequence):
        return key + '={' + ','.join(map(str, value)) + '}'
    else:
        return key + '=' + str(value)


def _format_range(value) -> str:
    if isinstance(value, str):
        return normalize_range_name(value)
    elif isinstance(value, Range):
        begin, end = value.first, value.last
    else:
        begin, end = value
    begin, end = normalize_range_name((str(begin), str(end)))
    return begin + '/' + end


def format_cmdpar(cmd, key: str, value) -> str:
    """
    Format a single MAD-X command parameter.

    :param cmd: A MAD-X Command instance for which an argument is to be formatted
    :param key: name of the parameter
    :param value: argument value
    """
    key = _fix_name(str(key).lower())
    cmdpar = cmd.cmdpar[key]
    dtype = cmdpar.dtype
    # the empty string was used in earlier versions in place of None:
    if value is None or value == '':
        return u''

    # NUMERIC
    if dtype == PARAM_TYPE_LOGICAL:
        if isinstance(value, bool):         return key + '=' + str(value).lower()
    if dtype in (PARAM_TYPE_LOGICAL,
                 PARAM_TYPE_INTEGER,
                 PARAM_TYPE_DOUBLE,
                 PARAM_TYPE_CONSTRAINT,
                 # NOTE: allow passing scalar values to numeric arrays, mainly
                 # useful because many of the `match` command parameters are
                 # arrays, but we usually call it with a single sequence and
                 # would like to treat it similar to the `twiss` command:
                 PARAM_TYPE_LOGICAL_ARRAY,
                 PARAM_TYPE_INTEGER_ARRAY,
                 PARAM_TYPE_DOUBLE_ARRAY):
        if isinstance(value, bool):         return key + '=' + str(int(value))
        if isinstance(value, Number):       return key + '=' + str(value)
        if isinstance(value, str):          return key + ':=' + value
    if dtype == PARAM_TYPE_CONSTRAINT:
        if isinstance(value, Constraint):
            constr = []
            if value.min is not None:
                constr.append(key + '>' + str(value.min))
            if value.max is not None:
                constr.append(key + '<' + str(value.max))
            if constr:
                return u', '.join(constr)
            else:
                return key + '=' + str(value.val)
    if dtype in (PARAM_TYPE_LOGICAL_ARRAY,
                 PARAM_TYPE_INTEGER_ARRAY,
                 PARAM_TYPE_DOUBLE_ARRAY):
        if isinstance(value, abc.Sequence):
            if all(isinstance(v, Number) for v in value):
                return key + '={' + ','.join(map(str, value)) + '}'
            else:
                return key + ':={' + ','.join(map(str, value)) + '}'

    # STRING
    def format_str(value):
        if key in QUOTED_PARAMS:
            return mad_quote(value)
        # NOTE: MAD-X stops parsing the current argument as soon as it
        # encounters another parameter name of the current command:
        elif is_identifier(value) and value not in cmd:
            return value
        else:
            return mad_quote(value.lower())
    if dtype == PARAM_TYPE_STRING:
        if key == 'range' or isinstance(value, Range):
            return key + '=' + _format_range(value)
        if isinstance(value, str):
            return key + '=' + format_str(value)
    if dtype == PARAM_TYPE_STRING_ARRAY:
        # NOTE: allowing single scalar value to STRING_ARRAY, mainly useful
        # for `match`, see above.
        if key == 'range' or isinstance(value, Range):
            if isinstance(value, list):
                return key + '={' + ','.join(map(_format_range, value)) + '}'
            return key + '=' + _format_range(value)
        if isinstance(value, str):
            return key + '=' + format_str(value)
        if isinstance(value, abc.Sequence):
            return key + '={' + ','.join(map(format_str, value)) + '}'

    raise TypeError('Unexpected command argument type: {}={!r} ({})'
                    .format(key, value, type(value)))


def format_command(*args, **kwargs) -> str:
    """
    Create a MAD-X command from its name and parameter list.

    :param cmd: base command (serves as template for parameter types)
    :param args: initial bareword command arguments (including command name!)
    :param kwargs: following named command arguments
    :returns: command string

    Examples:

    >>> format_command('twiss', sequence='lhc')
    'twiss, sequence=lhc;'

    >>> format_command('option', echo=True)
    'option, echo;'

    >>> format_command('constraint', betx=Constraint(max=3.13))
    'constraint, betx<3.13;'
    """
    cmd, args = args[0], args[1:]
    if isinstance(cmd, str):
        _args = [cmd] + list(args)
        _keys = ordered_keys(kwargs)
        _args += [format_param(k, kwargs[k]) for k in _keys]
    else:
        _args = [cmd.name] + list(args)
        _keys = ordered_keys(kwargs)
        _args += [format_cmdpar(cmd, k, kwargs[k]) for k in _keys]
    return u', '.join(filter(None, _args)) + ';'


# validation of MAD-X expressions

class T(Enum):
    """Terminal/token type."""
    WHITESPACE = 0
    LPAREN     = 1
    RPAREN     = 2
    COMMA      = 3
    SIGN       = 4
    OPERATOR   = 5
    SYMBOL     = 6
    NUMBER     = 7
    END        = 8

    __str__ = __repr__ = lambda self: self.name


class N(Enum):
    """Nonterminal symbol."""
    start            = 0
    expression       = 1
    expression_inner = 2
    expression_tail  = 3
    symbol_tail      = 4
    argument_list    = 5
    argument_tail    = 6

    __str__ = __repr__ = lambda self: self.name


grammar = {
    N.start: [
        [N.expression, T.END],
    ],
    N.expression: [
        [T.WHITESPACE, N.expression],
        [N.expression_inner],
    ],
    N.expression_inner: [
        [T.SIGN, N.expression],
        [T.LPAREN, N.expression, T.RPAREN, N.expression_tail],
        [T.NUMBER, N.expression_tail],
        [T.SYMBOL, N.symbol_tail],
    ],
    N.expression_tail: [
        [T.WHITESPACE, N.expression_tail],
        [T.SIGN, N.expression],
        [T.OPERATOR, N.expression],
        [],
    ],
    N.symbol_tail: [
        [T.WHITESPACE, N.symbol_tail],
        [T.LPAREN, N.argument_list, T.RPAREN, N.expression_tail],
        [T.SIGN, N.expression],
        [T.OPERATOR, N.expression],
        [],
    ],
    N.argument_list: [
        [T.WHITESPACE, N.argument_list],
        [N.expression_inner, N.argument_tail],
        [],
    ],
    N.argument_tail: [
        [T.COMMA, N.argument_list],
        [],
    ],
}


def _regex(expr: str) -> callable:
    regex = re.compile(expr)
    def match(text, i):
        m = regex.match(text[i:])
        return m.end() if m else 0
    return match


def _choice(choices: str) -> callable:
    def match(text, i):
        return 1 if text[i] in choices else 0
    return match


_expr_tokens = {
    T.WHITESPACE:   _choice(' \t'),
    T.LPAREN:       _choice('('),
    T.RPAREN:       _choice(')'),
    T.COMMA:        _choice(','),
    T.SIGN:         _choice('+-'),
    T.OPERATOR:     _choice('/*^'),
    T.SYMBOL:       _regex(r'[a-zA-Z_][a-zA-Z0-9_.]*(->[a-zA-Z_][a-zA-Z0-9_]*)?'),
    T.NUMBER:       _regex(r'(\d+(\.\d*)?|\.\d+)([eE][+\-]?\d+)?'),
}

_expr_parser = Parser(T, grammar, N.start)


class Token(namedtuple('Token', ['type', 'start', 'length', 'expr'])):

    @property
    def text(self):
        return self.expr[self.start:self.start + self.length]

    def __repr__(self):
        return '{}({!r})'.format(self.type, self.text)


def tokenize(tokens, expr: str):
    i = 0
    stop = len(expr)
    while i < stop:
        for toktype, tokmatch in tokens:
            l = tokmatch(expr, i)
            if l > 0:
                yield Token(toktype, i, l, expr)
                i += l
                break
        else:
            raise ValueError("Unknown token {!r} at {!r}"
                             .format(expr[i], expr[:i+1]))


def check_expression(expr: str):
    """
    Check if the given expression is a valid MAD-X expression that is safe to
    pass to :meth:`cpymad.madx.Madx.eval`.

    :param expr:
    :returns: True
    :raises ValueError: if the expression is ill-formed

    Note that this function only recognizes a sane subset of the expressions
    accepted by MAD-X and rejects valid but strange ones such as a number
    formatting '.' representing zero.
    """
    expr = expr.strip().lower()
    tokens = list(tokenize(list(_expr_tokens.items()), expr))
    tokens.append(Token(T.END, len(expr), 0, expr))
    _expr_parser.parse(tokens)  # raises ValueError
    return True


# misc

@contextmanager
def temp_filename():
    """Get filename for use within 'with' block and delete file afterwards."""
    fd, filename = tempfile.mkstemp()
    os.close(fd)
    try:
        yield filename
    finally:
        try:
            os.remove(filename)
        except OSError:
            pass


class ChangeDirectory:

    """
    Context manager for temporarily changing current working directory.

    :param str path: new path name
    :param _os: module with ``getcwd`` and ``chdir`` functions
    """

    # Note that the code is generic enough to be applied for any temporary
    # value patch, but we currently only need change directory, and therefore
    # have named it so.

    def __init__(self, path, chdir, getcwd):
        self._chdir = chdir
        self._getcwd = getcwd
        # Contrary to common implementations of a similar context manager,
        # we change the path immediately in the constructor. That enables
        # this utility to be used without any 'with' statement:
        if path:
            self._restore = getcwd()
            chdir(path)
        else:
            self._restore = None

    def __enter__(self):
        """Enter 'with' context."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit 'with' context and restore old path."""
        if self._restore:
            self._chdir(self._restore)


@np.vectorize
def remove_count_suffix_from_name(name):
    """Return the :N suffix from an element name."""
    return name.rsplit(':', 1)[0]

1	"""
2	Utility functions used in other parts of the pymad package.
3	"""
4	import re	10✔
5	import os	10✔
6	import sys	10✔
7	import tempfile	10✔
8	from collections import namedtuple	10✔
9	from contextlib import contextmanager	10✔
10	from enum import Enum	10✔
11	from numbers import Number	10✔
12	import collections.abc as abc	10✔
13
14	import numpy as np	10✔
15
16	from cpymad.parsing import Parser	10✔
17	from cpymad.types import (	10✔
18	Range, Constraint,
19	PARAM_TYPE_LOGICAL, PARAM_TYPE_INTEGER,
20	PARAM_TYPE_DOUBLE, PARAM_TYPE_STRING, PARAM_TYPE_CONSTRAINT,
21	PARAM_TYPE_LOGICAL_ARRAY, PARAM_TYPE_INTEGER_ARRAY,
22	PARAM_TYPE_DOUBLE_ARRAY, PARAM_TYPE_STRING_ARRAY)
23
24
25	__all__ = [	10✔
26	'mad_quote',
27	'is_identifier',
28	'name_from_internal',
29	'name_to_internal',
30	'format_param',
31	'format_cmdpar',
32	'format_command',
33	'check_expression',
34	'temp_filename',
35	'ChangeDirectory',
36	]
37
38
39	# In CPython 3.6 dicts preserve insertion order (until deleting an element)
40	# Although, this is considered an implementation detail that should not be
41	# relied upon, we do so anyway:
42	ordered_keys = dict.keys if sys.version_info >= (3, 6) else sorted	10✔
43
44
45	def mad_quote(value: str) -> str:	10✔
46	"""Add quotes to a string value."""
47	if '"' not in value:	10✔
48	return '"' + value + '"'	10✔
49	if "'" not in value:	10✔
50	return "'" + value + "'"	10✔
51	# MAD-X doesn't do any unescaping (otherwise I'd simply use `json.dumps`):
52	raise ValueError("MAD-X unable to parse string with escaped quotes: {!r}"	10✔
53	.format(value))
54
55
56	def _fix_name(name: str) -> str:	10✔
57	if name.startswith('_'):	10✔
58	raise AttributeError("Unknown item: {!r}! Did you mean {!r}?"	10✔
59	.format(name, name.strip('_') + '_'))
60	if name.endswith('_'):	10✔
61	name = name[:-1]	10✔
62	return name	10✔
63
64
65	# precompile regexes for performance:
66	re_compile = lambda s: re.compile(s, re.IGNORECASE)	10✔
67	_re_is_identifier = re_compile(r'^[a-z_][a-z0-9_]*$')	10✔
68	_re_symbol = re_compile(r'([a-z_][a-z0-9._](->[a-z_][a-z0-9._](\[[0-9]+\])?)?)')	10✔
69	_re_element_internal = re_compile(r'^([a-z_][a-z0-9_.$]*)(:\d+)?$')	10✔
70	_re_element_external = re_compile(r'^([a-z_][a-z0-9_.$]*)(\[\d+\])?$')	10✔
71
72
73	def is_identifier(name: str) -> bool:	10✔
74	"""Check if ``name`` is a valid identifier in MAD-X."""
75	return bool(_re_is_identifier.match(name))	10✔
76
77
78	def expr_symbols(expr: str) -> set:	10✔
79	"""
80	Return all symbols names used in an expression.
81
82	For now this includes not only variables but also element attributes (e.g.
83	``quad->k1``) as well as function names (e.g. ``sin``).
84	"""
85	return {m[0] for m in _re_symbol.findall(expr)}	10✔
86
87
88	def name_from_internal(element_name: str) -> str:	10✔
89	"""
90	Convert element name from internal representation to user API. Example:
91
92	>>> name_from_internal("foo:1")
93	foo
94	>>> name_from_internal("foo:2")
95	foo[2]
96
97	Element names are stored with a ":d" suffix by MAD-X internally (data in
98	node/sequence structs), but users must use the syntax "elem[d]" to access
99	the corresponding elements. This function is used to transform any string
100	coming from the user before passing it to MAD-X.
101	"""
102	try:	10✔
103	name, count = _re_element_internal.match(element_name).groups()	10✔
104	except AttributeError:	10✔
105	raise ValueError("Not a valid MAD-X element name: {!r}"	10✔
106	.format(element_name)) from None
107	if count is None or count == ':1':	10✔
108	return name	10✔
109	return name + '[' + count[1:] + ']'	10✔
110
111
112	def _parse_element_name(element_name: str) -> tuple:	10✔
113	"""
114	Parse element name from user API. Example:
115
116	>>> _parse_element_name("foo")
117	(foo, None)
118	>>> _parse_element_name("foo[2]")
119	(foo, 2)
120
121	See :func:`name_from_internal' for further information.
122	"""
123	try:	10✔
124	name, count = _re_element_external.match(element_name).groups()	10✔
125	except AttributeError:	10✔
126	raise ValueError("Not a valid MAD-X element name: {!r}"	10✔
127	.format(element_name)) from None
128	if count is None:	10✔
129	return name, None	10✔
130	return name, int(count[1:-1])	10✔
131
132
133	def name_to_internal(element_name: str) -> str:	10✔
134	"""
135	Convert element name from user API to internal representation. Example:
136
137	>>> name_to_external("foo")
138	foo:1
139	>>> name_to_external("foo[2]")
140	foo:2
141
142	See :func:`name_from_internal` for further information.
143	"""
144	name, count = _parse_element_name(element_name)	10✔
145	return name + ':' + str(1 if count is None else count)	10✔
146
147
148	def normalize_range_name(name: str, elems=None) -> str:	10✔
149	"""Make element name usable as argument to the RANGE attribute."""
150	if isinstance(name, tuple):	10✔
151	return tuple(map(normalize_range_name, name))	10✔
152	if '/' in name:	10✔
153	return '/'.join(map(normalize_range_name, name.split('/')))	10✔
154	name = name.lower()	10✔
155	if name.endswith('$end') or name.endswith('$start'):	10✔
156	if elems is None:	10✔
157	return u'#s' if name.endswith('$start') else u'#e'	10✔
158	else:
159	return u'#s' if elems.index(name) == 0 else u'#e'	×
160	return name	10✔
161
162
163	QUOTED_PARAMS = {'file', 'halofile', 'sectorfile', 'trueprofile'	10✔
164	'pipefile', 'trackfile', 'summary_file', 'filename',
165	'echo', 'title', 'text', 'format', 'dir'}
166
167
168	def format_param(key: str, value) -> str:	10✔
169	"""
170	Format a single MAD-X command parameter.
171
172	This is the old version that does not use type information from MAD-X. It
173	is therefore not limited to existing MAD-X commands and attributes, but
174	also less reliable for producing valid MAD-X statements.
175	"""
176	if value is None:	10✔
177	return None	10✔
178	key = _fix_name(str(key).lower())	10✔
179	if isinstance(value, Constraint):	10✔
180	constr = []	10✔
181	if value.min is not None:	10✔
182	constr.append(key + '>' + str(value.min))	10✔
183	if value.max is not None:	10✔
184	constr.append(key + '<' + str(value.max))	10✔
185	if constr:	10✔
186	return u', '.join(constr)	10✔
187	else:
188	return key + '=' + str(value.val)	10✔
189	elif isinstance(value, bool):	10✔
190	return key + '=' + str(value).lower()	10✔
191	elif key == 'range' or isinstance(value, Range):	10✔
192	return key + '=' + _format_range(value)	10✔
193	# check for basestrings before abc.Sequence, because every
194	# string is also a Sequence:
195	elif isinstance(value, str):	10✔
196	if key in QUOTED_PARAMS:	10✔
197	return key + '=' + mad_quote(value)	10✔
198	else:
199	# MAD-X parses strings incorrectly, if followed by a boolean.
200	# E.g.: "beam, sequence=s1, -radiate;" does NOT work! Therefore,
201	# these values need to be quoted. (NOTE: MAD-X uses lower-case
202	# internally and the quotes prevent automatic case conversion)
203	return key + '=' + mad_quote(value.lower())	10✔
204	# don't quote expressions:
205	elif isinstance(value, str):	10✔
206	return key + ':=' + value	×
207	elif isinstance(value, abc.Sequence):	10✔
208	return key + '={' + ','.join(map(str, value)) + '}'	10✔
209	else:
210	return key + '=' + str(value)	10✔
211
212
213	def _format_range(value) -> str:	10✔
214	if isinstance(value, str):	10✔
215	return normalize_range_name(value)	10✔
216	elif isinstance(value, Range):	10✔
217	begin, end = value.first, value.last	10✔
218	else:
219	begin, end = value	10✔
220	begin, end = normalize_range_name((str(begin), str(end)))	10✔
221	return begin + '/' + end	10✔
222
223
224	def format_cmdpar(cmd, key: str, value) -> str:	10✔
225	"""
226	Format a single MAD-X command parameter.
227
228	:param cmd: A MAD-X Command instance for which an argument is to be formatted
229	:param key: name of the parameter
230	:param value: argument value
231	"""
232	key = _fix_name(str(key).lower())	10✔
233	cmdpar = cmd.cmdpar[key]	10✔
234	dtype = cmdpar.dtype	10✔
235	# the empty string was used in earlier versions in place of None:
236	if value is None or value == '':	10✔
237	return u''	10✔
238
239	# NUMERIC
240	if dtype == PARAM_TYPE_LOGICAL:	10✔
241	if isinstance(value, bool): return key + '=' + str(value).lower()	10✔
242	if dtype in (PARAM_TYPE_LOGICAL,	10✔
243	PARAM_TYPE_INTEGER,
244	PARAM_TYPE_DOUBLE,
245	PARAM_TYPE_CONSTRAINT,
246	# NOTE: allow passing scalar values to numeric arrays, mainly
247	# useful because many of the `match` command parameters are
248	# arrays, but we usually call it with a single sequence and
249	# would like to treat it similar to the `twiss` command:
250	PARAM_TYPE_LOGICAL_ARRAY,
251	PARAM_TYPE_INTEGER_ARRAY,
252	PARAM_TYPE_DOUBLE_ARRAY):
253	if isinstance(value, bool): return key + '=' + str(int(value))	10✔
254	if isinstance(value, Number): return key + '=' + str(value)	10✔
255	if isinstance(value, str): return key + ':=' + value	10✔
256	if dtype == PARAM_TYPE_CONSTRAINT:	10✔
257	if isinstance(value, Constraint):	10✔
258	constr = []	10✔
259	if value.min is not None:	10✔
260	constr.append(key + '>' + str(value.min))	10✔
261	if value.max is not None:	10✔
262	constr.append(key + '<' + str(value.max))	10✔
263	if constr:	10✔
264	return u', '.join(constr)	10✔
265	else:
266	return key + '=' + str(value.val)	10✔
267	if dtype in (PARAM_TYPE_LOGICAL_ARRAY,	10✔
268	PARAM_TYPE_INTEGER_ARRAY,
269	PARAM_TYPE_DOUBLE_ARRAY):
270	if isinstance(value, abc.Sequence):	10✔
271	if all(isinstance(v, Number) for v in value):	10✔
272	return key + '={' + ','.join(map(str, value)) + '}'	10✔
273	else:
274	return key + ':={' + ','.join(map(str, value)) + '}'	10✔
275
276	# STRING
277	def format_str(value):	10✔
278	if key in QUOTED_PARAMS:	10✔
279	return mad_quote(value)	10✔
280	# NOTE: MAD-X stops parsing the current argument as soon as it
281	# encounters another parameter name of the current command:
282	elif is_identifier(value) and value not in cmd:	10✔
283	return value	10✔
284	else:
285	return mad_quote(value.lower())	10✔
286	if dtype == PARAM_TYPE_STRING:	10✔
287	if key == 'range' or isinstance(value, Range):	10✔
288	return key + '=' + _format_range(value)	10✔
289	if isinstance(value, str):	10✔
290	return key + '=' + format_str(value)	10✔
291	if dtype == PARAM_TYPE_STRING_ARRAY:	10✔
292	# NOTE: allowing single scalar value to STRING_ARRAY, mainly useful
293	# for `match`, see above.
294	if key == 'range' or isinstance(value, Range):	10✔
295	if isinstance(value, list):	10✔
296	return key + '={' + ','.join(map(_format_range, value)) + '}'	10✔
297	return key + '=' + _format_range(value)	10✔
298	if isinstance(value, str):	10✔
299	return key + '=' + format_str(value)	10✔
300	if isinstance(value, abc.Sequence):	10✔
301	return key + '={' + ','.join(map(format_str, value)) + '}'	10✔
302
303	raise TypeError('Unexpected command argument type: {}={!r} ({})'	×
304	.format(key, value, type(value)))
305
306
307	def format_command(args, *kwargs) -> str:	10✔
308	"""
309	Create a MAD-X command from its name and parameter list.
310
311	:param cmd: base command (serves as template for parameter types)
312	:param args: initial bareword command arguments (including command name!)
313	:param kwargs: following named command arguments
314	:returns: command string
315
316	Examples:
317
318	>>> format_command('twiss', sequence='lhc')
319	'twiss, sequence=lhc;'
320
321	>>> format_command('option', echo=True)
322	'option, echo;'
323
324	>>> format_command('constraint', betx=Constraint(max=3.13))
325	'constraint, betx<3.13;'
326	"""
327	cmd, args = args[0], args[1:]	10✔
328	if isinstance(cmd, str):	10✔
329	_args = [cmd] + list(args)	10✔
330	_keys = ordered_keys(kwargs)	10✔
331	_args += [format_param(k, kwargs[k]) for k in _keys]	10✔
332	else:
333	_args = [cmd.name] + list(args)	10✔
334	_keys = ordered_keys(kwargs)	10✔
335	_args += [format_cmdpar(cmd, k, kwargs[k]) for k in _keys]	10✔
336	return u', '.join(filter(None, _args)) + ';'	10✔
337
338
339	# validation of MAD-X expressions
340
341	class T(Enum):	10✔
342	"""Terminal/token type."""
343	WHITESPACE = 0	10✔
344	LPAREN = 1	10✔
345	RPAREN = 2	10✔
346	COMMA = 3	10✔
347	SIGN = 4	10✔
348	OPERATOR = 5	10✔
349	SYMBOL = 6	10✔
350	NUMBER = 7	10✔
351	END = 8	10✔
352
353	__str__ = __repr__ = lambda self: self.name	10✔
354
355
356	class N(Enum):	10✔
357	"""Nonterminal symbol."""
358	start = 0	10✔
359	expression = 1	10✔
360	expression_inner = 2	10✔
361	expression_tail = 3	10✔
362	symbol_tail = 4	10✔
363	argument_list = 5	10✔
364	argument_tail = 6	10✔
365
366	__str__ = __repr__ = lambda self: self.name	10✔
367
368
369	grammar = {	10✔
370	N.start: [
371	[N.expression, T.END],
372	],
373	N.expression: [
374	[T.WHITESPACE, N.expression],
375	[N.expression_inner],
376	],
377	N.expression_inner: [
378	[T.SIGN, N.expression],
379	[T.LPAREN, N.expression, T.RPAREN, N.expression_tail],
380	[T.NUMBER, N.expression_tail],
381	[T.SYMBOL, N.symbol_tail],
382	],
383	N.expression_tail: [
384	[T.WHITESPACE, N.expression_tail],
385	[T.SIGN, N.expression],
386	[T.OPERATOR, N.expression],
387	[],
388	],
389	N.symbol_tail: [
390	[T.WHITESPACE, N.symbol_tail],
391	[T.LPAREN, N.argument_list, T.RPAREN, N.expression_tail],
392	[T.SIGN, N.expression],
393	[T.OPERATOR, N.expression],
394	[],
395	],
396	N.argument_list: [
397	[T.WHITESPACE, N.argument_list],
398	[N.expression_inner, N.argument_tail],
399	[],
400	],
401	N.argument_tail: [
402	[T.COMMA, N.argument_list],
403	[],
404	],
405	}
406
407
408	def _regex(expr: str) -> callable:	10✔
409	regex = re.compile(expr)	10✔
410	def match(text, i):	10✔
411	m = regex.match(text[i:])	10✔
412	return m.end() if m else 0	10✔
413	return match	10✔
414
415
416	def _choice(choices: str) -> callable:	10✔
417	def match(text, i):	10✔
418	return 1 if text[i] in choices else 0	10✔
419	return match	10✔
420
421
422	_expr_tokens = {	10✔
423	T.WHITESPACE: _choice(' \t'),
424	T.LPAREN: _choice('('),
425	T.RPAREN: _choice(')'),
426	T.COMMA: _choice(','),
427	T.SIGN: _choice('+-'),
428	T.OPERATOR: _choice('/*^'),
429	T.SYMBOL: _regex(r'[a-zA-Z_][a-zA-Z0-9_.](->[a-zA-Z_][a-zA-Z0-9_])?'),
430	T.NUMBER: _regex(r'(\d+(\.\d*)?\|\.\d+)([eE][+\-]?\d+)?'),
431	}
432
433	_expr_parser = Parser(T, grammar, N.start)	10✔
434
435
436	class Token(namedtuple('Token', ['type', 'start', 'length', 'expr'])):	10✔
437
438	@property	10✔
439	def text(self):	6✔
440	return self.expr[self.start:self.start + self.length]	×
441
442	def __repr__(self):	10✔
443	return '{}({!r})'.format(self.type, self.text)	×
444
445
446	def tokenize(tokens, expr: str):	10✔
447	i = 0	10✔
448	stop = len(expr)	10✔
449	while i < stop:	10✔
450	for toktype, tokmatch in tokens:	10✔
451	l = tokmatch(expr, i)	10✔
452	if l > 0:	10✔
453	yield Token(toktype, i, l, expr)	10✔
454	i += l	10✔
455	break	10✔
456	else:
457	raise ValueError("Unknown token {!r} at {!r}"	10✔
458	.format(expr[i], expr[:i+1]))
459
460
461	def check_expression(expr: str):	10✔
462	"""
463	Check if the given expression is a valid MAD-X expression that is safe to
464	pass to :meth:`cpymad.madx.Madx.eval`.
465
466	:param expr:
467	:returns: True
468	:raises ValueError: if the expression is ill-formed
469
470	Note that this function only recognizes a sane subset of the expressions
471	accepted by MAD-X and rejects valid but strange ones such as a number
472	formatting '.' representing zero.
473	"""
474	expr = expr.strip().lower()	10✔
475	tokens = list(tokenize(list(_expr_tokens.items()), expr))	10✔
476	tokens.append(Token(T.END, len(expr), 0, expr))	10✔
477	_expr_parser.parse(tokens) # raises ValueError	10✔
478	return True	10✔
479
480
481	# misc
482
483	@contextmanager	10✔
484	def temp_filename():	6✔
485	"""Get filename for use within 'with' block and delete file afterwards."""
486	fd, filename = tempfile.mkstemp()	10✔
487	os.close(fd)	10✔
488	try:	10✔
489	yield filename	10✔
490	finally:
491	try:	10✔
492	os.remove(filename)	10✔
493	except OSError:	×
494	pass	×
495
496
497	class ChangeDirectory:	10✔
498
499	"""
500	Context manager for temporarily changing current working directory.
501
502	:param str path: new path name
503	:param _os: module with ``getcwd`` and ``chdir`` functions
504	"""
505
506	# Note that the code is generic enough to be applied for any temporary
507	# value patch, but we currently only need change directory, and therefore
508	# have named it so.
509
510	def __init__(self, path, chdir, getcwd):	10✔
511	self._chdir = chdir	10✔
512	self._getcwd = getcwd	10✔
513	# Contrary to common implementations of a similar context manager,
514	# we change the path immediately in the constructor. That enables
515	# this utility to be used without any 'with' statement:
516	if path:	10✔
517	self._restore = getcwd()	10✔
518	chdir(path)	10✔
519	else:
520	self._restore = None	×
521
522	def __enter__(self):	10✔
523	"""Enter 'with' context."""
524	return self	10✔
525
526	def __exit__(self, exc_type, exc_val, exc_tb):	10✔
527	"""Exit 'with' context and restore old path."""
528	if self._restore:	10✔
529	self._chdir(self._restore)	10✔
530
531
532	@np.vectorize	10✔
533	def remove_count_suffix_from_name(name):	6✔
534	"""Return the :N suffix from an element name."""
535	return name.rsplit(':', 1)[0]	10✔

hibtc / cpymad / 3676917429

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous