Merge pull request #121 from nneveu/cleanup

Cleanup

lcls_tools/common/devices/profile_monitor/profile_monitor.py → lcls_tools/common/devices/profile_monitor.py

@@ -1,7 +1,7 @@
 #!/usr/local/lcls/package/python/current/bin/python
 
 from epics import PV
-import lcls_tools.common.devices.profile_monitor.profmon_constants as pc
+import lcls_tools.common.devices.profmon_constants as pc
 from time import sleep
 from threading import Thread
 from numpy import array_equal

lcls_tools/common/devices/reader.py

@@ -1,14 +1,19 @@
 import os
 import yaml
-from typing import Union
+from typing import Union, Optional, Any, Dict
 from pydantic import ValidationError
+from lcls_tools.common.devices.screen import Screen, ScreenCollection
 from lcls_tools.common.devices.magnet import Magnet, MagnetCollection
 
 DEFAULT_YAML_LOCATION = "./lcls_tools/common/devices/yaml/"
 
 
-def _find_yaml_file(area: str) -> str:
-    filename = area + ".yaml"
+def _find_yaml_file(area: str = None, beampath: Optional[str] = None) -> str:
+    if area:
+        filename = area + ".yaml"
+    if beampath:
+        filename = "beam_paths.yaml"
+
     path = os.path.join(DEFAULT_YAML_LOCATION, filename)
     if os.path.isfile(path):
         return os.path.abspath(path)
@@ -18,32 +23,70 @@ def _find_yaml_file(area: str) -> str:
         )
 
 
-def create_magnet(
-    area: str = None, name: str = None
-) -> Union[None, Magnet, MagnetCollection]:
+def _device_data(
+    area: str = None,
+) -> Union[None, Dict[str, Any]]:
     if area:
         try:
             location = _find_yaml_file(
                 area=area,
             )
             with open(location, "r") as device_file:
-                config_data = yaml.safe_load(device_file)
-                if name:
-                    magnet_data = config_data["magnets"][name]
-                    # this data is not available from YAML directly in this form, so we add it here.
-                    magnet_data.update({"name": name})
-                    return Magnet(**magnet_data)
-                else:
-                    return MagnetCollection(**config_data)
+                device_data = yaml.safe_load(device_file)
+                return device_data
+
         except FileNotFoundError:
             print(f"Could not find yaml file for area: {area}")
             return None
+
+    else:
+        print("Please provide a machine area to create a magnet from.")
+        return None
+
+
+def create_beampath():
+    raise NotImplementedError
+
+
+def create_magnet(
+    area: str = None, name: str = None
+) -> Union[None, Magnet, MagnetCollection]:
+    device_data = _device_data(area=area)
+    if not device_data:
+        return None
+
+    if name:
+        try:
+            magnet_data = device_data["magnets"][name]
+            # this data is not available from YAML directly in this form, so we add it here.
+            magnet_data.update({"name": name})
+            return Magnet(**magnet_data)
         except KeyError:
-            print(f"Could not find name {name} in file for area: {area}")
-            return None
+            print(f"Magnet {name} does not exist in {area}.")
         except ValidationError as field_error:
             print(field_error)
             return None
     else:
-        print("Please provide a machine area to create a magnet from.")
+        return MagnetCollection(**device_data)
+
+
+def create_screen(
+    area: str = None, name: str = None
+) -> Union[None, Screen, ScreenCollection]:
+    device_data = _device_data(area=area)
+    if not device_data:
         return None
+
+    if name:
+        try:
+            screen_data = device_data["screens"][name]
+            # this data is not available from YAML directly in this form, so we add it here.
+            screen_data.update({"name": name})
+            return Screen(**screen_data)
+        except KeyError:
+            print(f"Screen {name} does not exist in {area}.")
+        except ValidationError as field_error:
+            print(field_error)
+            return None
+    else:
+        return ScreenCollection(**device_data)

lcls_tools/common/devices/yaml/generate.py

@@ -196,8 +196,8 @@ def extract_screens(self, area: Union[str, List[str]] = ["HTR"]):
             "IMAGE": "image",
             "Image:ArrayData": "image",
             "RESOLUTION": None,
-            "Image:ArraySizeX_RBV": "n_row",
-            "Image:ArraySizeY_RBV": "n_col",
+            "Image:ArraySize0_RBV": "n_row",
+            "Image:ArraySize1_RBV": "n_col",
             "N_OF_COL": "n_col",
             "N_OF_ROW": "n_row",
             "N_OF_BITS": "n_bits",

tests/unit_tests/lcls_tools/common/devices/profile_monitor/test_profile_monitor.py → tests/unit_tests/lcls_tools/common/devices/test_profile_monitor.py

@@ -3,8 +3,8 @@
 import inspect
 
 # Local imports
-import lcls_tools.common.devices.profile_monitor.profmon_constants as pc
-from lcls_tools.common.devices.profile_monitor.profile_monitor import (
+import lcls_tools.common.devices.profmon_constants as pc
+from lcls_tools.common.devices.profile_monitor import (
     ProfMon,
     get_profile_monitors,
 )

tests/unit_tests/lcls_tools/common/devices/test_reader.py

@@ -10,6 +10,10 @@ def setUp(self) -> None:
         self.config_location = "./tests/datasets/devices/config/magnet/"
         self.typical_config = os.path.join(self.config_location, "typical_magnet.yaml")
         self.bad_config = os.path.join(self.config_location, "bad_magnet.yaml")
+        # set up patch so that each magnet is constructured with ALL ctrl options
+        self.ctrl_options_patch = patch("epics.PV.get_ctrlvars", new_callable=MagicMock)
+        self.mock_ctrl_options = self.ctrl_options_patch.start()
+        self.mock_ctrl_options.return_value = {"enum_strs": tuple("READY")}
         return super().setUp()
 
     def test_bad_file_location_raises_when_finding(self):

tests/unit_tests/lcls_tools/common/devices/screen/test_screen.py → tests/unit_tests/lcls_tools/common/devices/test_screen.py

@@ -3,7 +3,7 @@
 import time
 import unittest
 from unittest.mock import PropertyMock, patch
-from lcls_tools.common.devices.screen.reader import create_screen
+from lcls_tools.common.devices.reader import create_screen
 import h5py
 import numpy as np
 
@@ -15,11 +15,11 @@ def setUp(self) -> None:
         return super().setUp()
 
     @patch(
-        "lcls_tools.common.devices.screen.screen.Screen.image",
+        "lcls_tools.common.devices.screen.Screen.image",
         new_callable=PropertyMock,
     )
     @patch(
-        "lcls_tools.common.devices.screen.screen.Screen.image_timestamp",
+        "lcls_tools.common.devices.screen.Screen.image_timestamp",
         new_callable=PropertyMock,
     )
     def test_hdf5_saving(self, mock_timestamp, mock_image):
@@ -54,11 +54,11 @@ def test_hdf5_saving(self, mock_timestamp, mock_image):
                     )
 
     @patch(
-        "lcls_tools.common.devices.screen.screen.Screen.image",
+        "lcls_tools.common.devices.screen.Screen.image",
         new_callable=PropertyMock,
     )
     @patch(
-        "lcls_tools.common.devices.screen.screen.Screen.image_timestamp",
+        "lcls_tools.common.devices.screen.Screen.image_timestamp",
         new_callable=PropertyMock,
     )
     def test_hdf5_with_user_metadata(self, mock_timestamp, mock_image):