access and query beam list (#124)

Author: rdemaria

src/cpymad/clibmadx.pxd

@@ -238,6 +238,7 @@ cdef extern from "madX/mad_gvar.h" nogil:
     command_list* defined_commands  # with base types, but no user elements
     int start_var               # start of variables after predefined constants
     int_array* deco             # temporary buffer for polished expressions
+    command_list* beam_list     # beam list
 
 
 # Function declarations:

src/cpymad/libmadx.pyx

@@ -122,6 +122,10 @@ __all__ = [
     'get_defined_command',
     'get_defined_command_names',
 
+    # defined commands
+    'get_beam',
+    'get_beam_names',
+
     # imported from 'os' for convenience in madx.Madx and should not be
     # considered part of the public interface:
     'getcwd',
@@ -911,6 +915,17 @@ def get_defined_command_names() -> list:
     """Return list of MAD-X command names."""
     return _name_list(clib.defined_commands.list)
 
+def get_beam(beam_name: str) -> dict:
+    """Return MAD-X beam as dict of values."""
+    cdef bytes _beam_name = _cstr(beam_name)
+    cdef int index = clib.name_list_pos(_beam_name, clib.beam_list.list)
+    if index == -1:
+        raise ValueError("Invalid beam: {!r}".format(beam_name))
+    return _parse_command(clib.beam_list.commands[index])
+
+def get_beam_names() -> list:
+    """Return list of MAD-X beam names."""
+    return _name_list(clib.beam_list.list)
 
 def is_sequence_expanded(sequence_name: str) -> bool:
     """

src/cpymad/madx.py

@@ -221,6 +221,7 @@ def __init__(self, libmadx=None, command_log=None, stdout=None,
         self.elements = GlobalElementList(self)
         self.base_types = BaseTypeMap(self)
         self.sequence = SequenceMap(self)
+        self.beams = BeamMap(self)
         self.table = TableMap(self._libmadx)
         self._enter_count = 0
         self._batch = None
@@ -606,6 +607,29 @@ def __call__(self):
             return None
 
 
+class BeamMap(_Mapping):
+    """Mapping of all beams (:class:`Beam`) in memory."""
+
+    def __init__(self, madx):
+        self._madx = madx
+        self._libmadx = madx._libmadx
+
+    def __iter__(self):
+        return iter(self._libmadx.get_beam_names())
+
+    def __getitem__(self, name):
+        try:
+            return Command(self._madx, self._libmadx.get_beam(name))
+        except ValueError:
+            raise KeyError("Unknown beam: {!r}".format(name)) from None
+
+    def __contains__(self, name):
+        return name in self._libmadx.get_beam_names()
+
+    def __len__(self):
+        return len(self._libmadx.get_beam_names())
+
+
 class TableMap(_Mapping):
 
     """Mapping of all tables (:class:`Table`) in memory."""

test/test_madx.py

@@ -1032,3 +1032,9 @@ def test_dframe_after_use(mad):
     # this line. It does not represent desired behaviour!
     assert mad.table.twiss.row_names() == \
         ['#s', '#e', 'dfd', 'mqd', 'dff', 'mqf']
+
+
+def test_beam_list(mad, lib):
+    mad.beam(pc=100)
+    assert lib.get_beam('default_beam')['data']['pc'].value == 100
+    assert lib.get_beam_names() == ['default_beam']