From 946f0e72df074813c7958cf8552e15dccb3dca78 Mon Sep 17 00:00:00 2001 From: Matt McDermott Date: Wed, 13 Mar 2024 10:53:25 -0700 Subject: [PATCH 1/9] add improved issue templates --- .github/ISSUE_TEMPLATE/bug_report.yaml | 37 +++++++++++++++++++++ .github/ISSUE_TEMPLATE/feature_request.yaml | 30 +++++++++++++++++ 2 files changed, 67 insertions(+) create mode 100644 .github/ISSUE_TEMPLATE/bug_report.yaml create mode 100644 .github/ISSUE_TEMPLATE/feature_request.yaml diff --git a/.github/ISSUE_TEMPLATE/bug_report.yaml b/.github/ISSUE_TEMPLATE/bug_report.yaml new file mode 100644 index 00000000..96770b3e --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.yaml @@ -0,0 +1,37 @@ +name: Report a Bug +description: Use this template to report an alab_os bug. +labels: ["bug"] +title: "[Bug]: " +body: + - type: textarea + id: current-behavior + attributes: + label: Current behavior + description: What bad behavior do you see? + validations: + required: true + + - type: textarea + id: expected-behavior + attributes: + label: Expected Behavior + description: What did you expect to see? + validations: + required: true + + - type: textarea + id: code-snippet + attributes: + label: Minimal example + description: If possible, provide a code snippet relevant to reproducing this bug. + render: Python + validations: + required: false + + - type: textarea + id: files + attributes: + label: Relevant files/images/logs + description: Please upload relevant files to help reproduce this bug, or logs if helpful. + validations: + required: false diff --git a/.github/ISSUE_TEMPLATE/feature_request.yaml b/.github/ISSUE_TEMPLATE/feature_request.yaml new file mode 100644 index 00000000..076b1afb --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.yaml @@ -0,0 +1,30 @@ +name: Feature Request +description: Use this template to request a new feature in alab_os. +labels: ["feature"] +title: "[Feature Request]: " +body: + - type: textarea + id: feature + attributes: + label: Feature Requested + description: Specify the feature and provide examples or use cases. + validations: + required: true + + - type: textarea + id: solution + attributes: + label: Proposed Solution + description: Share your thoughts on how the feature could be implemented. + placeholder: Implement a new method that ... + validations: + required: true + + - type: textarea + id: relevant + attributes: + label: Relevant Information + description: Additional context or links for understanding or implementing the feature. + placeholder: Use cases, related discussions, relevant literature, ... + validations: + required: false From 02fa5f6609753eb243e023ef49c7a22a66604998 Mon Sep 17 00:00:00 2001 From: Matt McDermott Date: Wed, 13 Mar 2024 16:56:19 -0700 Subject: [PATCH 2/9] add sample_view.update_sample_metadata(), plus a test --- alab_management/sample_view/sample_view.py | 79 ++++------ tests/test_sample_view.py | 163 ++++++--------------- 2 files changed, 78 insertions(+), 164 deletions(-) diff --git a/alab_management/sample_view/sample_view.py b/alab_management/sample_view/sample_view.py index 1ec32397..74c30aa1 100644 --- a/alab_management/sample_view/sample_view.py +++ b/alab_management/sample_view/sample_view.py @@ -103,8 +103,7 @@ def add_sample_positions_to_db( if re.search(r"[$.]", name) is not None: raise ValueError( - f"Unsupported sample position name: {name}. " - f"Sample position name should not contain '.' or '$'" + f"Unsupported sample position name: {name}. " f"Sample position name should not contain '.' or '$'" ) sample_pos_ = self._sample_positions_collection.find_one({"name": name}) @@ -146,9 +145,7 @@ def request_sample_positions( for sample_position in sample_positions ] - if len(sample_positions_request) != len( - {sample_position.prefix for sample_position in sample_positions_request} - ): + if len(sample_positions_request) != len({sample_position.prefix for sample_position in sample_positions_request}): raise ValueError("Duplicated sample_positions in one request.") # check if there are enough positions @@ -165,15 +162,13 @@ def request_sample_positions( with self._lock(): # pylint: disable=not-callable available_positions: Dict[str, List[Dict[str, Union[str, bool]]]] = {} for sample_position in sample_positions_request: - result = self.get_available_sample_position( - task_id, position_prefix=sample_position.prefix - ) + result = self.get_available_sample_position(task_id, position_prefix=sample_position.prefix) if not result or len(result) < sample_position.number: return None # we try to choose the position that has already been locked by this task - available_positions[sample_position.prefix] = sorted( - result, key=lambda task: int(task["need_release"]) - )[: sample_position.number] + available_positions[sample_position.prefix] = sorted(result, key=lambda task: int(task["need_release"]))[ + : sample_position.number + ] return available_positions def get_sample_position(self, position: str) -> Optional[Dict[str, Any]]: @@ -184,9 +179,7 @@ def get_sample_position(self, position: str) -> Optional[Dict[str, Any]]: """ return self._sample_positions_collection.find_one({"name": position}) - def get_sample_position_status( - self, position: str - ) -> Tuple[SamplePositionStatus, Optional[ObjectId]]: + def get_sample_position_status(self, position: str) -> Tuple[SamplePositionStatus, Optional[ObjectId]]: """ Get the status of a sample position. @@ -225,9 +218,7 @@ def get_sample_position_parent_device(self, position: str) -> Optional[str]: for position="furnace_1/tray" will properly return "furnace_1" even if "furnace_1/tray/1" and "furnace_1/tray/2" are in the database _as long as "furnace_1" is the parent device of both!_). """ - sample_positions = self._sample_positions_collection.find( - {"name": {"$regex": f"^{position}"}} - ) + sample_positions = self._sample_positions_collection.find({"name": {"$regex": f"^{position}"}}) parent_devices = list({sp.get("parent_device") for sp in sample_positions}) if len(parent_devices) == 0: raise ValueError(f"No sample position(s) beginning with: {position}") @@ -240,14 +231,9 @@ def get_sample_position_parent_device(self, position: str) -> Optional[str]: def is_unoccupied_position(self, position: str) -> bool: """Tell if a sample position is unoccupied or not.""" - return ( - self.get_sample_position_status(position)[0] - is not SamplePositionStatus.OCCUPIED - ) + return self.get_sample_position_status(position)[0] is not SamplePositionStatus.OCCUPIED - def get_available_sample_position( - self, task_id: ObjectId, position_prefix: str - ) -> List[Dict[str, Union[str, bool]]]: + def get_available_sample_position(self, task_id: ObjectId, position_prefix: str) -> List[Dict[str, Union[str, bool]]]: """ Check if the position is occupied. @@ -255,12 +241,7 @@ def get_available_sample_position( The entry need_release indicates whether a sample position needs to be released when __exit__ method is called in the ``SamplePositionsLock``. """ - if ( - self._sample_positions_collection.find_one( - {"name": {"$regex": f"^{re.escape(position_prefix)}"}} - ) - is None - ): + if self._sample_positions_collection.find_one({"name": {"$regex": f"^{re.escape(position_prefix)}"}}) is None: raise ValueError(f"Cannot find device with prefix: {position_prefix}") available_sample_positions = self._sample_positions_collection.find( @@ -278,9 +259,7 @@ def get_available_sample_position( ) available_sp_names = [] for sample_position in available_sample_positions: - status, current_task_id = self.get_sample_position_status( - sample_position["name"] - ) + status, current_task_id = self.get_sample_position_status(sample_position["name"]) if status is SamplePositionStatus.EMPTY or task_id == current_task_id: available_sp_names.append( { @@ -298,9 +277,7 @@ def lock_sample_position(self, task_id: ObjectId, position: str): if sample_status is SamplePositionStatus.OCCUPIED: raise ValueError(f"Position ({position}) is currently occupied") if sample_status is SamplePositionStatus.LOCKED: - raise ValueError( - f"Position is currently locked by task: {current_task_id}" - ) + raise ValueError(f"Position is currently locked by task: {current_task_id}") self._sample_positions_collection.update_one( {"name": position}, @@ -327,12 +304,7 @@ def release_sample_position(self, position: str): def get_sample_positions_by_task(self, task_id: Optional[ObjectId]) -> List[str]: """Get the list of sample positions that is locked by a task (given task id).""" - return [ - sample_position["name"] - for sample_position in self._sample_positions_collection.find( - {"task_id": task_id} - ) - ] + return [sample_position["name"] for sample_position in self._sample_positions_collection.find({"task_id": task_id})] ################################################################# # operations related to samples # @@ -355,10 +327,7 @@ def create_sample( raise ValueError(f"Requested position ({position}) is not EMPTY.") if re.search(r"[.$]", name) is not None: - raise ValueError( - f"Unsupported sample name: {name}. " - f"Sample name should not contain '.' or '$'" - ) + raise ValueError(f"Unsupported sample name: {name}. " f"Sample name should not contain '.' or '$'") entry = { "name": name, @@ -424,6 +393,20 @@ def update_sample_task_id(self, sample_id: ObjectId, task_id: Optional[ObjectId] }, ) + def update_sample_metadata(self, sample_id: ObjectId, metadata: Dict[str, Any]): + """Update the metadata for a sample. This adds new metadata or updates existing metadata.""" + result = self._sample_collection.find_one({"_id": sample_id}) + if result is None: + raise ValueError(f"Cannot find sample with id: {sample_id}") + + update_dict = {f"metadata.{k}": v for k, v in metadata.items()} + update_dict["last_updated"] = datetime.now() + + self._sample_collection.update_one( + {"_id": sample_id}, + {"$set": update_dict}, + ) + def move_sample(self, sample_id: ObjectId, position: Optional[str]): """Update the sample with new position.""" result = self._sample_collection.find_one({"_id": sample_id}) @@ -434,9 +417,7 @@ def move_sample(self, sample_id: ObjectId, position: Optional[str]): return if position is not None and not self.is_unoccupied_position(position): - raise ValueError( - f"Requested position ({position}) is not EMPTY or LOCKED by other task." - ) + raise ValueError(f"Requested position ({position}) is not EMPTY or LOCKED by other task.") self._sample_collection.update_one( {"_id": sample_id}, diff --git a/tests/test_sample_view.py b/tests/test_sample_view.py index 5fae1d6e..36cb2747 100644 --- a/tests/test_sample_view.py +++ b/tests/test_sample_view.py @@ -9,9 +9,7 @@ from alab_management.scripts.setup_lab import setup_lab -def occupy_sample_positions( - sample_positions, sample_view: SampleView, task_id: ObjectId -): +def occupy_sample_positions(sample_positions, sample_view: SampleView, task_id: ObjectId): for sample_positions_ in sample_positions.values(): for sample_position_ in sample_positions_: sample_view.lock_sample_position(task_id, sample_position_["name"]) @@ -48,29 +46,18 @@ def tearDown(self) -> None: self.sample_view._sample_collection.drop() @contextmanager - def request_sample_positions( - self, sample_positions_list, task_id: ObjectId, _timeout=None - ): + def request_sample_positions(self, sample_positions_list, task_id: ObjectId, _timeout=None): cnt = 0 - sample_positions = self.sample_view.request_sample_positions( - task_id=task_id, sample_positions=sample_positions_list - ) - while ( - _timeout is not None and sample_positions is None and _timeout >= cnt / 10 - ): - sample_positions = self.sample_view.request_sample_positions( - task_id=task_id, sample_positions=sample_positions_list - ) + sample_positions = self.sample_view.request_sample_positions(task_id=task_id, sample_positions=sample_positions_list) + while _timeout is not None and sample_positions is None and _timeout >= cnt / 10: + sample_positions = self.sample_view.request_sample_positions(task_id=task_id, sample_positions=sample_positions_list) cnt += 1 time.sleep(0.1) if sample_positions is not None: occupy_sample_positions(sample_positions, self.sample_view, task_id) yield ( - { - prefix: [sp["name"] for sp in sps] - for prefix, sps in sample_positions.items() - } + {prefix: [sp["name"] for sp in sps] for prefix, sps in sample_positions.items()} if sample_positions is not None else None ) @@ -102,6 +89,16 @@ def test_get_sample(self): with self.assertRaises(ValueError): self.sample_view.get_sample(sample_id=ObjectId()) + def test_update_sample_metadata(self): + sample_id = self.sample_view.create_sample("test_sample", position=None, metadata={"test_param": "test_value"}) + self.sample_view.update_sample_metadata(sample_id=sample_id, metadata={"test_param2": "test_value2"}) + sample = self.sample_view.get_sample(sample_id=sample_id) + self.assertDictEqual({"test_param": "test_value", "test_param2": "test_value2"}, sample.metadata) + + # try to update a non-exist sample + with self.assertRaises(ValueError): + self.sample_view.update_sample_metadata(sample_id=ObjectId(), metadata={"test_param": "test_value"}) + def test_move_sample(self): sample_id = self.sample_view.create_sample("test", position=None) sample_id_2 = self.sample_view.create_sample("test", position=None) @@ -126,9 +123,7 @@ def test_move_sample(self): # try to move a sample to an occupied position with self.assertRaises(ValueError): - self.sample_view.move_sample( - sample_id=sample_id_2, position="furnace_table" - ) + self.sample_view.move_sample(sample_id=sample_id_2, position="furnace_table") sample_2 = self.sample_view.get_sample(sample_id=sample_id_2) self.assertEqual(None, sample_2.position) @@ -149,9 +144,7 @@ def test_lock_sample_position(self): "LOCKED", self.sample_view.get_sample_position_status("furnace_table")[0].name, ) - self.assertListEqual( - ["furnace_table"], self.sample_view.get_sample_positions_by_task(task_id) - ) + self.assertListEqual(["furnace_table"], self.sample_view.get_sample_positions_by_task(task_id)) self.sample_view.release_sample_position("furnace_table") @@ -159,9 +152,7 @@ def test_lock_sample_position(self): # try to lock a sample position that already has a sample with self.assertRaises(ValueError): - self.sample_view.lock_sample_position( - task_id=task_id, position="furnace_table" - ) + self.sample_view.lock_sample_position(task_id=task_id, position="furnace_table") # try to lock a sample position with a sample that has the same task id self.sample_view.update_sample_task_id(sample_id=sample_id, task_id=task_id) @@ -194,17 +185,13 @@ def test_lock_sample_position(self): def test_request_sample_position_single(self): task_id = ObjectId() - with self.request_sample_positions( - ["furnace_table", "furnace_1/inside"], task_id - ) as sample_positions: + with self.request_sample_positions(["furnace_table", "furnace_1/inside"], task_id) as sample_positions: self.assertFalse(sample_positions is None) for sample_position_prefix, sample_position in sample_positions.items(): self.assertTrue(sample_position[0].startswith(sample_position_prefix)) self.assertEqual( "LOCKED", - self.sample_view.get_sample_position_status(sample_position[0])[ - 0 - ].name, + self.sample_view.get_sample_position_status(sample_position[0])[0].name, ) self.assertEqual( task_id, @@ -216,21 +203,15 @@ def test_request_sample_position_single(self): "EMPTY", self.sample_view.get_sample_position_status(sample_position[0])[0].name, ) - self.assertEqual( - None, self.sample_view.get_sample_position_status(sample_position[0])[1] - ) + self.assertEqual(None, self.sample_view.get_sample_position_status(sample_position[0])[1]) def test_request_device_timeout(self): task_id = ObjectId() task_id_2 = ObjectId() - with self.request_sample_positions( - ["furnace_table", "furnace_1/inside"], task_id - ) as sample_positions: + with self.request_sample_positions(["furnace_table", "furnace_1/inside"], task_id) as sample_positions: self.assertFalse(sample_positions is None) - with self.request_sample_positions( - ["furnace_table", "furnace_1/inside"], task_id_2 - ) as _sample_positions: + with self.request_sample_positions(["furnace_table", "furnace_1/inside"], task_id_2) as _sample_positions: self.assertIs(None, _sample_positions) def test_request_sample_positions_twice(self): @@ -256,9 +237,7 @@ def test_request_sample_positions_twice(self): "LOCKED", self.sample_view.get_sample_position_status("furnace_temp/1")[0].name, ) - with self.request_sample_positions( - ["furnace_table", "furnace_temp/1"], task_id - ) as sample_positions_: + with self.request_sample_positions(["furnace_table", "furnace_temp/1"], task_id) as sample_positions_: self.assertDictEqual( { "furnace_table": ["furnace_table"], @@ -268,28 +247,18 @@ def test_request_sample_positions_twice(self): ) self.assertEqual( "LOCKED", - self.sample_view.get_sample_position_status("furnace_table")[ - 0 - ].name, + self.sample_view.get_sample_position_status("furnace_table")[0].name, ) - with self.request_sample_positions( - ["furnace_table"], task_id - ) as sample_positions__: - self.assertDictEqual( - {"furnace_table": ["furnace_table"]}, sample_positions__ - ) + with self.request_sample_positions(["furnace_table"], task_id) as sample_positions__: + self.assertDictEqual({"furnace_table": ["furnace_table"]}, sample_positions__) self.assertEqual( "LOCKED", - self.sample_view.get_sample_position_status("furnace_table")[ - 0 - ].name, + self.sample_view.get_sample_position_status("furnace_table")[0].name, ) self.assertEqual( "LOCKED", - self.sample_view.get_sample_position_status("furnace_table")[ - 0 - ].name, + self.sample_view.get_sample_position_status("furnace_table")[0].name, ) self.assertEqual( @@ -301,9 +270,7 @@ def test_request_sample_positions_twice(self): "EMPTY", self.sample_view.get_sample_position_status("furnace_table")[0].name, ) - self.assertEqual( - None, self.sample_view.get_sample_position_status("furnace_table")[1] - ) + self.assertEqual(None, self.sample_view.get_sample_position_status("furnace_table")[1]) def test_request_sample_positions_occupied(self): task_id = ObjectId() @@ -320,97 +287,63 @@ def test_request_sample_positions_occupied(self): self.sample_view.update_sample_task_id(sample_id, task_id) - with self.request_sample_positions( - ["furnace_table", "furnace_1/inside"], task_id - ): + with self.request_sample_positions(["furnace_table", "furnace_1/inside"], task_id): self.assertEqual( "OCCUPIED", self.sample_view.get_sample_position_status("furnace_table")[0].name, ) - self.assertEqual( - task_id, self.sample_view.get_sample_position_status("furnace_table")[1] - ) + self.assertEqual(task_id, self.sample_view.get_sample_position_status("furnace_table")[1]) self.assertEqual( "OCCUPIED", self.sample_view.get_sample_position_status("furnace_table")[0].name, ) - self.assertEqual( - None, self.sample_view.get_sample_position("furnace_table")["task_id"] - ) + self.assertEqual(None, self.sample_view.get_sample_position("furnace_table")["task_id"]) def test_request_multiple_sample_positions(self): task_id = ObjectId() for j in range(1, 5): - with self.request_sample_positions( - [{"prefix": "furnace_temp", "number": j}], task_id - ) as sample_positions: + with self.request_sample_positions([{"prefix": "furnace_temp", "number": j}], task_id) as sample_positions: self.assertFalse(sample_positions is None) for sample_position_prefix, sample_position in sample_positions.items(): for i in range(j): - self.assertTrue( - sample_position[i].startswith(sample_position_prefix) - ) + self.assertTrue(sample_position[i].startswith(sample_position_prefix)) self.assertEqual( "LOCKED", - self.sample_view.get_sample_position_status( - sample_position[i] - )[0].name, + self.sample_view.get_sample_position_status(sample_position[i])[0].name, ) self.assertEqual( task_id, - self.sample_view.get_sample_position_status( - sample_position[i] - )[1], + self.sample_view.get_sample_position_status(sample_position[i])[1], ) for sample_position in sample_positions.values(): for i in range(j): self.assertEqual( "EMPTY", - self.sample_view.get_sample_position_status(sample_position[i])[ - 0 - ].name, + self.sample_view.get_sample_position_status(sample_position[i])[0].name, ) self.assertEqual( None, - self.sample_view.get_sample_position_status(sample_position[i])[ - 1 - ], + self.sample_view.get_sample_position_status(sample_position[i])[1], ) # try when requesting sample positions more than we have in the lab - with self.assertRaises(ValueError), self.request_sample_positions( - [{"prefix": "furnace_temp", "number": 5}], task_id - ): + with self.assertRaises(ValueError), self.request_sample_positions([{"prefix": "furnace_temp", "number": 5}], task_id): pass def test_request_multiple_sample_positions_multiple_tasks(self): task_id_1 = ObjectId() task_id_2 = ObjectId() - with self.request_sample_positions( - [{"prefix": "furnace_temp", "number": 2}], task_id_1 - ) as sample_positions: + with self.request_sample_positions([{"prefix": "furnace_temp", "number": 2}], task_id_1) as sample_positions: self.assertEqual(2, len(sample_positions["furnace_temp"])) - self.assertTrue( - sample_positions["furnace_temp"][0].startswith("furnace_temp") - ) - self.assertTrue( - sample_positions["furnace_temp"][1].startswith("furnace_temp") - ) - with self.request_sample_positions( - [{"prefix": "furnace_temp", "number": 2}], task_id_2 - ) as sample_positions_: + self.assertTrue(sample_positions["furnace_temp"][0].startswith("furnace_temp")) + self.assertTrue(sample_positions["furnace_temp"][1].startswith("furnace_temp")) + with self.request_sample_positions([{"prefix": "furnace_temp", "number": 2}], task_id_2) as sample_positions_: self.assertEqual(2, len(sample_positions_["furnace_temp"])) - self.assertTrue( - sample_positions_["furnace_temp"][0].startswith("furnace_temp") - ) - self.assertTrue( - sample_positions_["furnace_temp"][1].startswith("furnace_temp") - ) - with self.request_sample_positions( - [{"prefix": "furnace_temp", "number": 4}], task_id_2 - ) as sample_positions_: + self.assertTrue(sample_positions_["furnace_temp"][0].startswith("furnace_temp")) + self.assertTrue(sample_positions_["furnace_temp"][1].startswith("furnace_temp")) + with self.request_sample_positions([{"prefix": "furnace_temp", "number": 4}], task_id_2) as sample_positions_: self.assertIs(None, sample_positions_) From 3342b4ce5939258d1029f2b2c0030c1417c3d21a Mon Sep 17 00:00:00 2001 From: Matt McDermott Date: Thu, 14 Mar 2024 14:13:09 -0700 Subject: [PATCH 3/9] add update_sample_metadata() function to lab_view for convenient access --- alab_management/lab_view.py | 106 +++++++++++++----------------------- 1 file changed, 39 insertions(+), 67 deletions(-) diff --git a/alab_management/lab_view.py b/alab_management/lab_view.py index 0eb45eee..4cf1ae96 100644 --- a/alab_management/lab_view.py +++ b/alab_management/lab_view.py @@ -47,13 +47,7 @@ def preprocess(cls, values): """Preprocess the request to make sure the request is in the correct format.""" values = values["__root__"] # if the sample position request is string, we will automatically add a number attribute = 1. - values = { - k: [ - SamplePositionRequest.from_str(v_) if isinstance(v_, str) else v_ - for v_ in v - ] - for k, v in values.items() - } + values = {k: [SamplePositionRequest.from_str(v_) if isinstance(v_, str) else v_ for v_ in v] for k, v in values.items()} return {"__root__": values} @@ -66,9 +60,7 @@ class LabView: def __init__(self, task_id: ObjectId): self._task_view = TaskView() - self.__task_entry = self._task_view.get_task( - task_id=task_id - ) # will throw error if task_id does not exist + self.__task_entry = self._task_view.get_task(task_id=task_id) # will throw error if task_id does not exist self._experiment_view = ExperimentView() self._task_id = task_id self._sample_view = SampleView() @@ -86,9 +78,7 @@ def task_id(self) -> ObjectId: @contextmanager def request_resources( self, - resource_request: Dict[ - Optional[Union[Type[BaseDevice], str]], Dict[str, Union[str, int]] - ], + resource_request: Dict[Optional[Union[Type[BaseDevice], str]], Dict[str, Union[str, int]]], priority: Optional[int] = None, timeout: Optional[float] = None, ): @@ -114,18 +104,13 @@ def request_resources( """ priority = priority or self.priority - self._task_view.update_status( - task_id=self.task_id, status=TaskStatus.REQUESTING_RESOURCES - ) - result = self._resource_requester.request_resources( - resource_request=resource_request, timeout=timeout, priority=priority - ) + self._task_view.update_status(task_id=self.task_id, status=TaskStatus.REQUESTING_RESOURCES) + result = self._resource_requester.request_resources(resource_request=resource_request, timeout=timeout, priority=priority) devices = result["devices"] sample_positions = result["sample_positions"] request_id = result["request_id"] devices = { - device_type: self._device_client.create_device_wrapper(device_name) - for device_type, device_name in devices.items() + device_type: self._device_client.create_device_wrapper(device_name) for device_type, device_name in devices.items() } # type: ignore self._task_view.update_status(task_id=self.task_id, status=TaskStatus.RUNNING) yield devices, sample_positions @@ -141,9 +126,7 @@ def _sample_name_to_id(self, sample_name: str) -> ObjectId: for sample in self.__task_entry["samples"]: if sample["name"] == sample_name: return sample["sample_id"] - raise ValueError( - f"No sample with name \"{sample_name}\" found for task \"{self.__task_entry['type']}\"" - ) + raise ValueError(f"No sample with name \"{sample_name}\" found for task \"{self.__task_entry['type']}\"") def get_sample(self, sample: Union[ObjectId, str]) -> Sample: """ @@ -169,23 +152,15 @@ def move_sample(self, sample: Union[ObjectId, str], position: Optional[str]): :py:meth:`move_sample ` """ # check if this sample position is locked by current task - if ( - position is not None - and self._sample_view.get_sample_position_status(position)[1] - != self._task_id - ): - raise ValueError( - f"Cannot move sample to the new sample position ({position}) without locking it." - ) + if position is not None and self._sample_view.get_sample_position_status(position)[1] != self._task_id: + raise ValueError(f"Cannot move sample to the new sample position ({position}) without locking it.") # check if this sample is owned by current task sample_entry = self.get_sample(sample=sample) if sample_entry.task_id != self._task_id: raise ValueError("Cannot move a sample that is not belong to this task.") - return self._sample_view.move_sample( - sample_id=sample_entry.sample_id, position=position - ) + return self._sample_view.move_sample(sample_id=sample_entry.sample_id, position=position) def get_locked_sample_positions(self) -> List[str]: """Get a list of sample positions that are occupied by this task.""" @@ -195,9 +170,25 @@ def get_sample_position_parent_device(self, position: str) -> Optional[str]: """Get the name of the device that owns the sample position.""" return self._sample_view.get_sample_position_parent_device(position=position) - def run_subtask( - self, task: Type[BaseTask], samples: List[Union[ObjectId, str]], **kwargs - ): + def update_sample_metadata(self, sample: Union[ObjectId, str], metadata: Dict[str, Any]): + """ + Update the metadata of a sample. `sample` can be given as either an ObjectId corresponding to sample_id, + or as a string corresponding to the sample's name within the experiment. + + See Also + -------- + :py:meth:`update_sample_metadata ` + """ + if isinstance(sample, str): + sample_id = self._sample_name_to_id(sample) + elif isinstance(sample, ObjectId): + sample_id = sample + else: + raise TypeError("sample must be a sample name (str) or id (ObjectId)") + + return self._sample_view.update_sample_metadata(sample_id=sample_id, metadata=metadata) + + def run_subtask(self, task: Type[BaseTask], samples: List[Union[ObjectId, str]], **kwargs): """Run a task as a subtask within the task. basically fills in task_id and lab_view for you. this command blocks until the subtask is completed. @@ -231,12 +222,8 @@ def run_subtask( **kwargs, ) except Exception as exc: - self._task_view.update_subtask_status( - task_id=task_id, subtask_id=subtask_id, status=TaskStatus.ERROR - ) - self._task_view.update_subtask_result( - task_id=task_id, subtask_id=subtask_id, result=str(exc) - ) + self._task_view.update_subtask_status(task_id=task_id, subtask_id=subtask_id, status=TaskStatus.ERROR) + self._task_view.update_subtask_result(task_id=task_id, subtask_id=subtask_id, result=str(exc)) raise Exception( "Failed to create subtask of type {} within task {} of type {}".format( task, @@ -254,17 +241,11 @@ def run_subtask( }, ) try: - self._task_view.update_subtask_status( - task_id=task_id, subtask_id=subtask_id, status=TaskStatus.RUNNING - ) + self._task_view.update_subtask_status(task_id=task_id, subtask_id=subtask_id, status=TaskStatus.RUNNING) result = subtask.run() # block until completion except Exception as exception: - self._task_view.update_subtask_status( - task_id=task_id, subtask_id=subtask_id, status=TaskStatus.ERROR - ) - self._task_view.update_subtask_result( - task_id=task_id, subtask_id=subtask_id, result=str(exception) - ) + self._task_view.update_subtask_status(task_id=task_id, subtask_id=subtask_id, status=TaskStatus.ERROR) + self._task_view.update_subtask_result(task_id=task_id, subtask_id=subtask_id, result=str(exception)) self.logger.system_log( level="ERROR", log_data={ @@ -278,12 +259,8 @@ def run_subtask( ) raise else: - self._task_view.update_subtask_status( - task_id=task_id, subtask_id=subtask_id, status=TaskStatus.COMPLETED - ) - self._task_view.update_subtask_result( - task_id=task_id, subtask_id=subtask_id, result=result - ) + self._task_view.update_subtask_status(task_id=task_id, subtask_id=subtask_id, status=TaskStatus.COMPLETED) + self._task_view.update_subtask_result(task_id=task_id, subtask_id=subtask_id, result=result) self.logger.system_log( level="INFO", log_data={ @@ -326,19 +303,14 @@ def update_result(self, name: str, value: Any): def request_cleanup(self): """Request cleanup of the task. This function will block until the task is cleaned up.""" - all_reserved_sample_positions = self._sample_view.get_sample_positions_by_task( - self.task_id - ) + all_reserved_sample_positions = self._sample_view.get_sample_positions_by_task(self.task_id) all_samples = self.__task_entry["samples"] all_positions_with_samples = [ - self._sample_view.get_sample(sample_entry["sample_id"]).position - for sample_entry in all_samples + self._sample_view.get_sample(sample_entry["sample_id"]).position for sample_entry in all_samples ] - all_positions_with_samples = [ - each for each in all_positions_with_samples if each - ] + all_positions_with_samples = [each for each in all_positions_with_samples if each] self.request_user_input( prompt="A unrecoverable error has occurred.\n" From 4322428d16a028a20692b4a2c2decbdcf8776970 Mon Sep 17 00:00:00 2001 From: Bernardus Rendy <37468936+bernardusrendy@users.noreply.github.com> Date: Sun, 17 Mar 2024 12:55:20 -0700 Subject: [PATCH 4/9] Python>=3.10, pylint, ruff, black --- alab_management/lab_view.py | 96 +++++++++--- alab_management/sample_view/sample_view.py | 65 ++++++-- alab_management/scripts/windows_service.py | 14 +- pyproject.toml | 4 +- tests/test_sample_view.py | 172 ++++++++++++++++----- 5 files changed, 261 insertions(+), 90 deletions(-) diff --git a/alab_management/lab_view.py b/alab_management/lab_view.py index 4cf1ae96..90f2e17a 100644 --- a/alab_management/lab_view.py +++ b/alab_management/lab_view.py @@ -47,7 +47,13 @@ def preprocess(cls, values): """Preprocess the request to make sure the request is in the correct format.""" values = values["__root__"] # if the sample position request is string, we will automatically add a number attribute = 1. - values = {k: [SamplePositionRequest.from_str(v_) if isinstance(v_, str) else v_ for v_ in v] for k, v in values.items()} + values = { + k: [ + SamplePositionRequest.from_str(v_) if isinstance(v_, str) else v_ + for v_ in v + ] + for k, v in values.items() + } return {"__root__": values} @@ -60,7 +66,9 @@ class LabView: def __init__(self, task_id: ObjectId): self._task_view = TaskView() - self.__task_entry = self._task_view.get_task(task_id=task_id) # will throw error if task_id does not exist + self.__task_entry = self._task_view.get_task( + task_id=task_id + ) # will throw error if task_id does not exist self._experiment_view = ExperimentView() self._task_id = task_id self._sample_view = SampleView() @@ -78,7 +86,9 @@ def task_id(self) -> ObjectId: @contextmanager def request_resources( self, - resource_request: Dict[Optional[Union[Type[BaseDevice], str]], Dict[str, Union[str, int]]], + resource_request: Dict[ + Optional[Union[Type[BaseDevice], str]], Dict[str, Union[str, int]] + ], priority: Optional[int] = None, timeout: Optional[float] = None, ): @@ -104,13 +114,18 @@ def request_resources( """ priority = priority or self.priority - self._task_view.update_status(task_id=self.task_id, status=TaskStatus.REQUESTING_RESOURCES) - result = self._resource_requester.request_resources(resource_request=resource_request, timeout=timeout, priority=priority) + self._task_view.update_status( + task_id=self.task_id, status=TaskStatus.REQUESTING_RESOURCES + ) + result = self._resource_requester.request_resources( + resource_request=resource_request, timeout=timeout, priority=priority + ) devices = result["devices"] sample_positions = result["sample_positions"] request_id = result["request_id"] devices = { - device_type: self._device_client.create_device_wrapper(device_name) for device_type, device_name in devices.items() + device_type: self._device_client.create_device_wrapper(device_name) + for device_type, device_name in devices.items() } # type: ignore self._task_view.update_status(task_id=self.task_id, status=TaskStatus.RUNNING) yield devices, sample_positions @@ -126,7 +141,9 @@ def _sample_name_to_id(self, sample_name: str) -> ObjectId: for sample in self.__task_entry["samples"]: if sample["name"] == sample_name: return sample["sample_id"] - raise ValueError(f"No sample with name \"{sample_name}\" found for task \"{self.__task_entry['type']}\"") + raise ValueError( + f"No sample with name \"{sample_name}\" found for task \"{self.__task_entry['type']}\"" + ) def get_sample(self, sample: Union[ObjectId, str]) -> Sample: """ @@ -152,15 +169,23 @@ def move_sample(self, sample: Union[ObjectId, str], position: Optional[str]): :py:meth:`move_sample ` """ # check if this sample position is locked by current task - if position is not None and self._sample_view.get_sample_position_status(position)[1] != self._task_id: - raise ValueError(f"Cannot move sample to the new sample position ({position}) without locking it.") + if ( + position is not None + and self._sample_view.get_sample_position_status(position)[1] + != self._task_id + ): + raise ValueError( + f"Cannot move sample to the new sample position ({position}) without locking it." + ) # check if this sample is owned by current task sample_entry = self.get_sample(sample=sample) if sample_entry.task_id != self._task_id: raise ValueError("Cannot move a sample that is not belong to this task.") - return self._sample_view.move_sample(sample_id=sample_entry.sample_id, position=position) + return self._sample_view.move_sample( + sample_id=sample_entry.sample_id, position=position + ) def get_locked_sample_positions(self) -> List[str]: """Get a list of sample positions that are occupied by this task.""" @@ -170,7 +195,9 @@ def get_sample_position_parent_device(self, position: str) -> Optional[str]: """Get the name of the device that owns the sample position.""" return self._sample_view.get_sample_position_parent_device(position=position) - def update_sample_metadata(self, sample: Union[ObjectId, str], metadata: Dict[str, Any]): + def update_sample_metadata( + self, sample: Union[ObjectId, str], metadata: Dict[str, Any] + ): """ Update the metadata of a sample. `sample` can be given as either an ObjectId corresponding to sample_id, or as a string corresponding to the sample's name within the experiment. @@ -186,9 +213,13 @@ def update_sample_metadata(self, sample: Union[ObjectId, str], metadata: Dict[st else: raise TypeError("sample must be a sample name (str) or id (ObjectId)") - return self._sample_view.update_sample_metadata(sample_id=sample_id, metadata=metadata) + return self._sample_view.update_sample_metadata( + sample_id=sample_id, metadata=metadata + ) - def run_subtask(self, task: Type[BaseTask], samples: List[Union[ObjectId, str]], **kwargs): + def run_subtask( + self, task: Type[BaseTask], samples: List[Union[ObjectId, str]], **kwargs + ): """Run a task as a subtask within the task. basically fills in task_id and lab_view for you. this command blocks until the subtask is completed. @@ -222,8 +253,12 @@ def run_subtask(self, task: Type[BaseTask], samples: List[Union[ObjectId, str]], **kwargs, ) except Exception as exc: - self._task_view.update_subtask_status(task_id=task_id, subtask_id=subtask_id, status=TaskStatus.ERROR) - self._task_view.update_subtask_result(task_id=task_id, subtask_id=subtask_id, result=str(exc)) + self._task_view.update_subtask_status( + task_id=task_id, subtask_id=subtask_id, status=TaskStatus.ERROR + ) + self._task_view.update_subtask_result( + task_id=task_id, subtask_id=subtask_id, result=str(exc) + ) raise Exception( "Failed to create subtask of type {} within task {} of type {}".format( task, @@ -241,11 +276,17 @@ def run_subtask(self, task: Type[BaseTask], samples: List[Union[ObjectId, str]], }, ) try: - self._task_view.update_subtask_status(task_id=task_id, subtask_id=subtask_id, status=TaskStatus.RUNNING) + self._task_view.update_subtask_status( + task_id=task_id, subtask_id=subtask_id, status=TaskStatus.RUNNING + ) result = subtask.run() # block until completion except Exception as exception: - self._task_view.update_subtask_status(task_id=task_id, subtask_id=subtask_id, status=TaskStatus.ERROR) - self._task_view.update_subtask_result(task_id=task_id, subtask_id=subtask_id, result=str(exception)) + self._task_view.update_subtask_status( + task_id=task_id, subtask_id=subtask_id, status=TaskStatus.ERROR + ) + self._task_view.update_subtask_result( + task_id=task_id, subtask_id=subtask_id, result=str(exception) + ) self.logger.system_log( level="ERROR", log_data={ @@ -259,8 +300,12 @@ def run_subtask(self, task: Type[BaseTask], samples: List[Union[ObjectId, str]], ) raise else: - self._task_view.update_subtask_status(task_id=task_id, subtask_id=subtask_id, status=TaskStatus.COMPLETED) - self._task_view.update_subtask_result(task_id=task_id, subtask_id=subtask_id, result=result) + self._task_view.update_subtask_status( + task_id=task_id, subtask_id=subtask_id, status=TaskStatus.COMPLETED + ) + self._task_view.update_subtask_result( + task_id=task_id, subtask_id=subtask_id, result=result + ) self.logger.system_log( level="INFO", log_data={ @@ -303,14 +348,19 @@ def update_result(self, name: str, value: Any): def request_cleanup(self): """Request cleanup of the task. This function will block until the task is cleaned up.""" - all_reserved_sample_positions = self._sample_view.get_sample_positions_by_task(self.task_id) + all_reserved_sample_positions = self._sample_view.get_sample_positions_by_task( + self.task_id + ) all_samples = self.__task_entry["samples"] all_positions_with_samples = [ - self._sample_view.get_sample(sample_entry["sample_id"]).position for sample_entry in all_samples + self._sample_view.get_sample(sample_entry["sample_id"]).position + for sample_entry in all_samples ] - all_positions_with_samples = [each for each in all_positions_with_samples if each] + all_positions_with_samples = [ + each for each in all_positions_with_samples if each + ] self.request_user_input( prompt="A unrecoverable error has occurred.\n" diff --git a/alab_management/sample_view/sample_view.py b/alab_management/sample_view/sample_view.py index 74c30aa1..9e029862 100644 --- a/alab_management/sample_view/sample_view.py +++ b/alab_management/sample_view/sample_view.py @@ -103,7 +103,8 @@ def add_sample_positions_to_db( if re.search(r"[$.]", name) is not None: raise ValueError( - f"Unsupported sample position name: {name}. " f"Sample position name should not contain '.' or '$'" + f"Unsupported sample position name: {name}. " + f"Sample position name should not contain '.' or '$'" ) sample_pos_ = self._sample_positions_collection.find_one({"name": name}) @@ -145,7 +146,9 @@ def request_sample_positions( for sample_position in sample_positions ] - if len(sample_positions_request) != len({sample_position.prefix for sample_position in sample_positions_request}): + if len(sample_positions_request) != len( + {sample_position.prefix for sample_position in sample_positions_request} + ): raise ValueError("Duplicated sample_positions in one request.") # check if there are enough positions @@ -162,13 +165,15 @@ def request_sample_positions( with self._lock(): # pylint: disable=not-callable available_positions: Dict[str, List[Dict[str, Union[str, bool]]]] = {} for sample_position in sample_positions_request: - result = self.get_available_sample_position(task_id, position_prefix=sample_position.prefix) + result = self.get_available_sample_position( + task_id, position_prefix=sample_position.prefix + ) if not result or len(result) < sample_position.number: return None # we try to choose the position that has already been locked by this task - available_positions[sample_position.prefix] = sorted(result, key=lambda task: int(task["need_release"]))[ - : sample_position.number - ] + available_positions[sample_position.prefix] = sorted( + result, key=lambda task: int(task["need_release"]) + )[: sample_position.number] return available_positions def get_sample_position(self, position: str) -> Optional[Dict[str, Any]]: @@ -179,7 +184,9 @@ def get_sample_position(self, position: str) -> Optional[Dict[str, Any]]: """ return self._sample_positions_collection.find_one({"name": position}) - def get_sample_position_status(self, position: str) -> Tuple[SamplePositionStatus, Optional[ObjectId]]: + def get_sample_position_status( + self, position: str + ) -> Tuple[SamplePositionStatus, Optional[ObjectId]]: """ Get the status of a sample position. @@ -218,7 +225,9 @@ def get_sample_position_parent_device(self, position: str) -> Optional[str]: for position="furnace_1/tray" will properly return "furnace_1" even if "furnace_1/tray/1" and "furnace_1/tray/2" are in the database _as long as "furnace_1" is the parent device of both!_). """ - sample_positions = self._sample_positions_collection.find({"name": {"$regex": f"^{position}"}}) + sample_positions = self._sample_positions_collection.find( + {"name": {"$regex": f"^{position}"}} + ) parent_devices = list({sp.get("parent_device") for sp in sample_positions}) if len(parent_devices) == 0: raise ValueError(f"No sample position(s) beginning with: {position}") @@ -231,9 +240,14 @@ def get_sample_position_parent_device(self, position: str) -> Optional[str]: def is_unoccupied_position(self, position: str) -> bool: """Tell if a sample position is unoccupied or not.""" - return self.get_sample_position_status(position)[0] is not SamplePositionStatus.OCCUPIED + return ( + self.get_sample_position_status(position)[0] + is not SamplePositionStatus.OCCUPIED + ) - def get_available_sample_position(self, task_id: ObjectId, position_prefix: str) -> List[Dict[str, Union[str, bool]]]: + def get_available_sample_position( + self, task_id: ObjectId, position_prefix: str + ) -> List[Dict[str, Union[str, bool]]]: """ Check if the position is occupied. @@ -241,7 +255,12 @@ def get_available_sample_position(self, task_id: ObjectId, position_prefix: str) The entry need_release indicates whether a sample position needs to be released when __exit__ method is called in the ``SamplePositionsLock``. """ - if self._sample_positions_collection.find_one({"name": {"$regex": f"^{re.escape(position_prefix)}"}}) is None: + if ( + self._sample_positions_collection.find_one( + {"name": {"$regex": f"^{re.escape(position_prefix)}"}} + ) + is None + ): raise ValueError(f"Cannot find device with prefix: {position_prefix}") available_sample_positions = self._sample_positions_collection.find( @@ -259,7 +278,9 @@ def get_available_sample_position(self, task_id: ObjectId, position_prefix: str) ) available_sp_names = [] for sample_position in available_sample_positions: - status, current_task_id = self.get_sample_position_status(sample_position["name"]) + status, current_task_id = self.get_sample_position_status( + sample_position["name"] + ) if status is SamplePositionStatus.EMPTY or task_id == current_task_id: available_sp_names.append( { @@ -277,7 +298,9 @@ def lock_sample_position(self, task_id: ObjectId, position: str): if sample_status is SamplePositionStatus.OCCUPIED: raise ValueError(f"Position ({position}) is currently occupied") if sample_status is SamplePositionStatus.LOCKED: - raise ValueError(f"Position is currently locked by task: {current_task_id}") + raise ValueError( + f"Position is currently locked by task: {current_task_id}" + ) self._sample_positions_collection.update_one( {"name": position}, @@ -304,7 +327,12 @@ def release_sample_position(self, position: str): def get_sample_positions_by_task(self, task_id: Optional[ObjectId]) -> List[str]: """Get the list of sample positions that is locked by a task (given task id).""" - return [sample_position["name"] for sample_position in self._sample_positions_collection.find({"task_id": task_id})] + return [ + sample_position["name"] + for sample_position in self._sample_positions_collection.find( + {"task_id": task_id} + ) + ] ################################################################# # operations related to samples # @@ -327,7 +355,10 @@ def create_sample( raise ValueError(f"Requested position ({position}) is not EMPTY.") if re.search(r"[.$]", name) is not None: - raise ValueError(f"Unsupported sample name: {name}. " f"Sample name should not contain '.' or '$'") + raise ValueError( + f"Unsupported sample name: {name}. " + f"Sample name should not contain '.' or '$'" + ) entry = { "name": name, @@ -417,7 +448,9 @@ def move_sample(self, sample_id: ObjectId, position: Optional[str]): return if position is not None and not self.is_unoccupied_position(position): - raise ValueError(f"Requested position ({position}) is not EMPTY or LOCKED by other task.") + raise ValueError( + f"Requested position ({position}) is not EMPTY or LOCKED by other task." + ) self._sample_collection.update_one( {"_id": sample_id}, diff --git a/alab_management/scripts/windows_service.py b/alab_management/scripts/windows_service.py index e9b45a45..9d8790b4 100644 --- a/alab_management/scripts/windows_service.py +++ b/alab_management/scripts/windows_service.py @@ -25,9 +25,10 @@ def process_launch_worker(): Launches the worker process. This is a workaround to redirect the output of the worker process to a file. """ - with open("D:\\alabos_service.log", "a", encoding="utf-8") as sys.stdout, open( - "D:\\alabos_service.log", "a", encoding="utf-8" - ) as sys.stderr: + with ( + open("D:\\alabos_service.log", "a", encoding="utf-8") as sys.stdout, + open("D:\\alabos_service.log", "a", encoding="utf-8") as sys.stderr, + ): launch_worker([]) @@ -68,9 +69,10 @@ def SvcDoRun(self): def main(self): """Main function of the service.""" - with redirect_stdout( - open("D:\\alabos_service.log", "a", encoding="utf-8") - ), redirect_stderr(open("D:\\alabos_service.log", "a", encoding="utf-8")): + with ( + redirect_stdout(open("D:\\alabos_service.log", "a", encoding="utf-8")), + redirect_stderr(open("D:\\alabos_service.log", "a", encoding="utf-8")), + ): self.running = True self.warned = False self.alabos_thread = Thread( diff --git a/pyproject.toml b/pyproject.toml index 5bccc1e6..83cd19af 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,13 +24,11 @@ classifiers = [ "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", "Topic :: Database :: Front-Ends", "Topic :: Other/Nonlisted Topic", "Topic :: Scientific/Engineering", ] -requires-python = ">=3.8.0" +requires-python = ">=3.10.0" dependencies = [ "toml>=0.10.1", "pymongo>=3.12.3", diff --git a/tests/test_sample_view.py b/tests/test_sample_view.py index 36cb2747..b1a1f3b3 100644 --- a/tests/test_sample_view.py +++ b/tests/test_sample_view.py @@ -9,7 +9,9 @@ from alab_management.scripts.setup_lab import setup_lab -def occupy_sample_positions(sample_positions, sample_view: SampleView, task_id: ObjectId): +def occupy_sample_positions( + sample_positions, sample_view: SampleView, task_id: ObjectId +): for sample_positions_ in sample_positions.values(): for sample_position_ in sample_positions_: sample_view.lock_sample_position(task_id, sample_position_["name"]) @@ -46,18 +48,29 @@ def tearDown(self) -> None: self.sample_view._sample_collection.drop() @contextmanager - def request_sample_positions(self, sample_positions_list, task_id: ObjectId, _timeout=None): + def request_sample_positions( + self, sample_positions_list, task_id: ObjectId, _timeout=None + ): cnt = 0 - sample_positions = self.sample_view.request_sample_positions(task_id=task_id, sample_positions=sample_positions_list) - while _timeout is not None and sample_positions is None and _timeout >= cnt / 10: - sample_positions = self.sample_view.request_sample_positions(task_id=task_id, sample_positions=sample_positions_list) + sample_positions = self.sample_view.request_sample_positions( + task_id=task_id, sample_positions=sample_positions_list + ) + while ( + _timeout is not None and sample_positions is None and _timeout >= cnt / 10 + ): + sample_positions = self.sample_view.request_sample_positions( + task_id=task_id, sample_positions=sample_positions_list + ) cnt += 1 time.sleep(0.1) if sample_positions is not None: occupy_sample_positions(sample_positions, self.sample_view, task_id) yield ( - {prefix: [sp["name"] for sp in sps] for prefix, sps in sample_positions.items()} + { + prefix: [sp["name"] for sp in sps] + for prefix, sps in sample_positions.items() + } if sample_positions is not None else None ) @@ -90,14 +103,22 @@ def test_get_sample(self): self.sample_view.get_sample(sample_id=ObjectId()) def test_update_sample_metadata(self): - sample_id = self.sample_view.create_sample("test_sample", position=None, metadata={"test_param": "test_value"}) - self.sample_view.update_sample_metadata(sample_id=sample_id, metadata={"test_param2": "test_value2"}) + sample_id = self.sample_view.create_sample( + "test_sample", position=None, metadata={"test_param": "test_value"} + ) + self.sample_view.update_sample_metadata( + sample_id=sample_id, metadata={"test_param2": "test_value2"} + ) sample = self.sample_view.get_sample(sample_id=sample_id) - self.assertDictEqual({"test_param": "test_value", "test_param2": "test_value2"}, sample.metadata) + self.assertDictEqual( + {"test_param": "test_value", "test_param2": "test_value2"}, sample.metadata + ) # try to update a non-exist sample with self.assertRaises(ValueError): - self.sample_view.update_sample_metadata(sample_id=ObjectId(), metadata={"test_param": "test_value"}) + self.sample_view.update_sample_metadata( + sample_id=ObjectId(), metadata={"test_param": "test_value"} + ) def test_move_sample(self): sample_id = self.sample_view.create_sample("test", position=None) @@ -123,7 +144,9 @@ def test_move_sample(self): # try to move a sample to an occupied position with self.assertRaises(ValueError): - self.sample_view.move_sample(sample_id=sample_id_2, position="furnace_table") + self.sample_view.move_sample( + sample_id=sample_id_2, position="furnace_table" + ) sample_2 = self.sample_view.get_sample(sample_id=sample_id_2) self.assertEqual(None, sample_2.position) @@ -144,7 +167,9 @@ def test_lock_sample_position(self): "LOCKED", self.sample_view.get_sample_position_status("furnace_table")[0].name, ) - self.assertListEqual(["furnace_table"], self.sample_view.get_sample_positions_by_task(task_id)) + self.assertListEqual( + ["furnace_table"], self.sample_view.get_sample_positions_by_task(task_id) + ) self.sample_view.release_sample_position("furnace_table") @@ -152,7 +177,9 @@ def test_lock_sample_position(self): # try to lock a sample position that already has a sample with self.assertRaises(ValueError): - self.sample_view.lock_sample_position(task_id=task_id, position="furnace_table") + self.sample_view.lock_sample_position( + task_id=task_id, position="furnace_table" + ) # try to lock a sample position with a sample that has the same task id self.sample_view.update_sample_task_id(sample_id=sample_id, task_id=task_id) @@ -185,13 +212,17 @@ def test_lock_sample_position(self): def test_request_sample_position_single(self): task_id = ObjectId() - with self.request_sample_positions(["furnace_table", "furnace_1/inside"], task_id) as sample_positions: + with self.request_sample_positions( + ["furnace_table", "furnace_1/inside"], task_id + ) as sample_positions: self.assertFalse(sample_positions is None) for sample_position_prefix, sample_position in sample_positions.items(): self.assertTrue(sample_position[0].startswith(sample_position_prefix)) self.assertEqual( "LOCKED", - self.sample_view.get_sample_position_status(sample_position[0])[0].name, + self.sample_view.get_sample_position_status(sample_position[0])[ + 0 + ].name, ) self.assertEqual( task_id, @@ -203,15 +234,21 @@ def test_request_sample_position_single(self): "EMPTY", self.sample_view.get_sample_position_status(sample_position[0])[0].name, ) - self.assertEqual(None, self.sample_view.get_sample_position_status(sample_position[0])[1]) + self.assertEqual( + None, self.sample_view.get_sample_position_status(sample_position[0])[1] + ) def test_request_device_timeout(self): task_id = ObjectId() task_id_2 = ObjectId() - with self.request_sample_positions(["furnace_table", "furnace_1/inside"], task_id) as sample_positions: + with self.request_sample_positions( + ["furnace_table", "furnace_1/inside"], task_id + ) as sample_positions: self.assertFalse(sample_positions is None) - with self.request_sample_positions(["furnace_table", "furnace_1/inside"], task_id_2) as _sample_positions: + with self.request_sample_positions( + ["furnace_table", "furnace_1/inside"], task_id_2 + ) as _sample_positions: self.assertIs(None, _sample_positions) def test_request_sample_positions_twice(self): @@ -237,7 +274,9 @@ def test_request_sample_positions_twice(self): "LOCKED", self.sample_view.get_sample_position_status("furnace_temp/1")[0].name, ) - with self.request_sample_positions(["furnace_table", "furnace_temp/1"], task_id) as sample_positions_: + with self.request_sample_positions( + ["furnace_table", "furnace_temp/1"], task_id + ) as sample_positions_: self.assertDictEqual( { "furnace_table": ["furnace_table"], @@ -247,18 +286,28 @@ def test_request_sample_positions_twice(self): ) self.assertEqual( "LOCKED", - self.sample_view.get_sample_position_status("furnace_table")[0].name, + self.sample_view.get_sample_position_status("furnace_table")[ + 0 + ].name, ) - with self.request_sample_positions(["furnace_table"], task_id) as sample_positions__: - self.assertDictEqual({"furnace_table": ["furnace_table"]}, sample_positions__) + with self.request_sample_positions( + ["furnace_table"], task_id + ) as sample_positions__: + self.assertDictEqual( + {"furnace_table": ["furnace_table"]}, sample_positions__ + ) self.assertEqual( "LOCKED", - self.sample_view.get_sample_position_status("furnace_table")[0].name, + self.sample_view.get_sample_position_status("furnace_table")[ + 0 + ].name, ) self.assertEqual( "LOCKED", - self.sample_view.get_sample_position_status("furnace_table")[0].name, + self.sample_view.get_sample_position_status("furnace_table")[ + 0 + ].name, ) self.assertEqual( @@ -270,7 +319,9 @@ def test_request_sample_positions_twice(self): "EMPTY", self.sample_view.get_sample_position_status("furnace_table")[0].name, ) - self.assertEqual(None, self.sample_view.get_sample_position_status("furnace_table")[1]) + self.assertEqual( + None, self.sample_view.get_sample_position_status("furnace_table")[1] + ) def test_request_sample_positions_occupied(self): task_id = ObjectId() @@ -287,63 +338,100 @@ def test_request_sample_positions_occupied(self): self.sample_view.update_sample_task_id(sample_id, task_id) - with self.request_sample_positions(["furnace_table", "furnace_1/inside"], task_id): + with self.request_sample_positions( + ["furnace_table", "furnace_1/inside"], task_id + ): self.assertEqual( "OCCUPIED", self.sample_view.get_sample_position_status("furnace_table")[0].name, ) - self.assertEqual(task_id, self.sample_view.get_sample_position_status("furnace_table")[1]) + self.assertEqual( + task_id, self.sample_view.get_sample_position_status("furnace_table")[1] + ) self.assertEqual( "OCCUPIED", self.sample_view.get_sample_position_status("furnace_table")[0].name, ) - self.assertEqual(None, self.sample_view.get_sample_position("furnace_table")["task_id"]) + self.assertEqual( + None, self.sample_view.get_sample_position("furnace_table")["task_id"] + ) def test_request_multiple_sample_positions(self): task_id = ObjectId() for j in range(1, 5): - with self.request_sample_positions([{"prefix": "furnace_temp", "number": j}], task_id) as sample_positions: + with self.request_sample_positions( + [{"prefix": "furnace_temp", "number": j}], task_id + ) as sample_positions: self.assertFalse(sample_positions is None) for sample_position_prefix, sample_position in sample_positions.items(): for i in range(j): - self.assertTrue(sample_position[i].startswith(sample_position_prefix)) + self.assertTrue( + sample_position[i].startswith(sample_position_prefix) + ) self.assertEqual( "LOCKED", - self.sample_view.get_sample_position_status(sample_position[i])[0].name, + self.sample_view.get_sample_position_status( + sample_position[i] + )[0].name, ) self.assertEqual( task_id, - self.sample_view.get_sample_position_status(sample_position[i])[1], + self.sample_view.get_sample_position_status( + sample_position[i] + )[1], ) for sample_position in sample_positions.values(): for i in range(j): self.assertEqual( "EMPTY", - self.sample_view.get_sample_position_status(sample_position[i])[0].name, + self.sample_view.get_sample_position_status(sample_position[i])[ + 0 + ].name, ) self.assertEqual( None, - self.sample_view.get_sample_position_status(sample_position[i])[1], + self.sample_view.get_sample_position_status(sample_position[i])[ + 1 + ], ) # try when requesting sample positions more than we have in the lab - with self.assertRaises(ValueError), self.request_sample_positions([{"prefix": "furnace_temp", "number": 5}], task_id): + with ( + self.assertRaises(ValueError), + self.request_sample_positions( + [{"prefix": "furnace_temp", "number": 5}], task_id + ), + ): pass def test_request_multiple_sample_positions_multiple_tasks(self): task_id_1 = ObjectId() task_id_2 = ObjectId() - with self.request_sample_positions([{"prefix": "furnace_temp", "number": 2}], task_id_1) as sample_positions: + with self.request_sample_positions( + [{"prefix": "furnace_temp", "number": 2}], task_id_1 + ) as sample_positions: self.assertEqual(2, len(sample_positions["furnace_temp"])) - self.assertTrue(sample_positions["furnace_temp"][0].startswith("furnace_temp")) - self.assertTrue(sample_positions["furnace_temp"][1].startswith("furnace_temp")) - with self.request_sample_positions([{"prefix": "furnace_temp", "number": 2}], task_id_2) as sample_positions_: + self.assertTrue( + sample_positions["furnace_temp"][0].startswith("furnace_temp") + ) + self.assertTrue( + sample_positions["furnace_temp"][1].startswith("furnace_temp") + ) + with self.request_sample_positions( + [{"prefix": "furnace_temp", "number": 2}], task_id_2 + ) as sample_positions_: self.assertEqual(2, len(sample_positions_["furnace_temp"])) - self.assertTrue(sample_positions_["furnace_temp"][0].startswith("furnace_temp")) - self.assertTrue(sample_positions_["furnace_temp"][1].startswith("furnace_temp")) - with self.request_sample_positions([{"prefix": "furnace_temp", "number": 4}], task_id_2) as sample_positions_: + self.assertTrue( + sample_positions_["furnace_temp"][0].startswith("furnace_temp") + ) + self.assertTrue( + sample_positions_["furnace_temp"][1].startswith("furnace_temp") + ) + with self.request_sample_positions( + [{"prefix": "furnace_temp", "number": 4}], task_id_2 + ) as sample_positions_: self.assertIs(None, sample_positions_) From 84b31674a6d22b2edc6867ff7475ca46e5deeb78 Mon Sep 17 00:00:00 2001 From: Bernardus Rendy <37468936+bernardusrendy@users.noreply.github.com> Date: Sun, 17 Mar 2024 13:43:18 -0700 Subject: [PATCH 5/9] Update black.yaml --- .github/workflows/black.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/black.yaml b/.github/workflows/black.yaml index ede5e898..52228e81 100644 --- a/.github/workflows/black.yaml +++ b/.github/workflows/black.yaml @@ -13,7 +13,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v2 with: - python-version: "3.9" # Specify the Python version you need + python-version: "3.10" # Specify the Python version you need - name: Install Black run: pip install black==22.3.0 - name: Run Black From 86b64d1940b7c84fc37ad20db71f58f5aec755f0 Mon Sep 17 00:00:00 2001 From: Bernardus Rendy <37468936+bernardusrendy@users.noreply.github.com> Date: Sun, 17 Mar 2024 13:50:50 -0700 Subject: [PATCH 6/9] black, CI, page.yaml --- .github/workflows/black.yaml | 2 +- .github/workflows/ci.yaml | 4 ++-- .github/workflows/page.yaml | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/black.yaml b/.github/workflows/black.yaml index 52228e81..02a9acd2 100644 --- a/.github/workflows/black.yaml +++ b/.github/workflows/black.yaml @@ -15,6 +15,6 @@ jobs: with: python-version: "3.10" # Specify the Python version you need - name: Install Black - run: pip install black==22.3.0 + run: pip install black - name: Run Black run: black ./alab_management diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 9395e10b..a2827af2 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -16,7 +16,7 @@ jobs: timeout-minutes: 60 strategy: matrix: - python-version: [ "3.8", "3.9", "3.10" ] + python-version: ["3.10"] services: rabbitmq: image: "rabbitmq:3.9" @@ -37,7 +37,7 @@ jobs: uses: actions/setup-python@v2 with: # python-version: ${{ matrix.python-version }} - python-version: "3.9" + python-version: "3.10" # cache: 'pip' # cache-dependency-path: - name: Set up environment diff --git a/.github/workflows/page.yaml b/.github/workflows/page.yaml index ac48179a..e6c71956 100644 --- a/.github/workflows/page.yaml +++ b/.github/workflows/page.yaml @@ -19,7 +19,7 @@ jobs: - name: Set up python uses: actions/setup-python@v2 with: - python-version: '3.8' + python-version: '3.10' - name: Set up dependencies run: pip install --quiet . - name: Compile sphinx From 1c2aa663d8dc69cbf496c4b00b9e8a1c8d513c92 Mon Sep 17 00:00:00 2001 From: Bernardus Rendy <37468936+bernardusrendy@users.noreply.github.com> Date: Sun, 17 Mar 2024 13:51:04 -0700 Subject: [PATCH 7/9] Ruff, Black, Lint to python 3.10 --- alab_management/builders/experimentbuilder.py | 20 ++++----- alab_management/builders/samplebuilder.py | 8 ++-- alab_management/builders/utils.py | 4 +- alab_management/config.py | 4 +- .../dashboard/routes/experiment.py | 6 +-- alab_management/device_manager.py | 13 +++--- alab_management/device_view/dbattributes.py | 6 +-- alab_management/device_view/device.py | 27 ++++++------ alab_management/device_view/device_view.py | 35 ++++++++-------- alab_management/experiment_manager.py | 8 ++-- .../completed_experiment_view.py | 6 +-- alab_management/experiment_view/experiment.py | 24 +++++------ .../experiment_view/experiment_view.py | 18 ++++---- alab_management/lab_view.py | 28 ++++++------- alab_management/logger.py | 25 ++++++----- .../sample_view/completed_sample_view.py | 4 +- alab_management/sample_view/sample.py | 14 +++---- alab_management/sample_view/sample_view.py | 42 +++++++++---------- .../task_manager/resource_requester.py | 42 +++++++++---------- alab_management/task_manager/task_manager.py | 18 ++++---- .../task_view/completed_task_view.py | 4 +- alab_management/task_view/task.py | 40 +++++++++--------- alab_management/task_view/task_view.py | 38 ++++++++--------- alab_management/user_input.py | 18 ++++---- alab_management/utils/data_objects.py | 13 +++--- alab_management/utils/db_lock.py | 5 +-- alab_management/utils/graph_ops.py | 8 ++-- alab_management/utils/module_ops.py | 5 +-- examples/fake_lab/tasks/ending.py | 4 +- examples/fake_lab/tasks/heating.py | 5 +-- examples/fake_lab/tasks/moving.py | 4 +- examples/fake_lab/tasks/starting.py | 3 +- pyproject.toml | 4 +- tests/fake_lab/tasks/ending.py | 4 +- tests/fake_lab/tasks/heating.py | 5 +-- tests/fake_lab/tasks/moving.py | 4 +- tests/fake_lab/tasks/starting.py | 3 +- 37 files changed, 247 insertions(+), 272 deletions(-) diff --git a/alab_management/builders/experimentbuilder.py b/alab_management/builders/experimentbuilder.py index c83c139b..e0cb4ceb 100644 --- a/alab_management/builders/experimentbuilder.py +++ b/alab_management/builders/experimentbuilder.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Any, Dict, List, Literal, Optional, Set, Union +from typing import Any, Literal from .samplebuilder import SampleBuilder @@ -13,20 +13,20 @@ class ExperimentBuilder: name (str): The name of the experiment. """ - def __init__(self, name: str, tags: Optional[List[str]] = None, **metadata): + def __init__(self, name: str, tags: list[str] | None = None, **metadata): """ Args: name (str): The name of the experiment. tags (List[str]): A list of tags to attach to the experiment. """ self.name = name - self._samples: List[SampleBuilder] = [] - self._tasks: Dict[str, Dict[str, Any]] = {} + self._samples: list[SampleBuilder] = [] + self._tasks: dict[str, dict[str, Any]] = {} self.tags = tags or [] self.metadata = metadata def add_sample( - self, name: str, tags: Optional[List[str]] = None, **metadata + self, name: str, tags: list[str] | None = None, **metadata ) -> SampleBuilder: """ Add a sample to the batch. Each sample already has multiple tasks binded to it. Each @@ -53,8 +53,8 @@ def add_task( self, task_id: str, task_name: str, - task_kwargs: Dict[str, Any], - samples: List[SampleBuilder], + task_kwargs: dict[str, Any], + samples: list[SampleBuilder], ) -> None: """ This function adds a task to the sample. You should use this function only for special cases which @@ -78,7 +78,7 @@ def add_task( "samples": [sample.name for sample in samples], } - def to_dict(self) -> Dict[str, Any]: + def to_dict(self) -> dict[str, Any]: """ Return a dictionary that can be used to generate an input file for the `experiment` to run. @@ -88,9 +88,9 @@ def to_dict(self) -> Dict[str, Any]: A dictionary that can be used to generate an input file for the `experiment` to run. """ - samples: List[Dict[str, Any]] = [] + samples: list[dict[str, Any]] = [] # tasks = [] - tasks: List[Dict[str, Union[str, Set[int], List]]] = [] + tasks: list[dict[str, str | set[int] | list]] = [] task_ids = {} for sample in self._samples: diff --git a/alab_management/builders/samplebuilder.py b/alab_management/builders/samplebuilder.py index ad84dd5b..93dd605c 100644 --- a/alab_management/builders/samplebuilder.py +++ b/alab_management/builders/samplebuilder.py @@ -1,6 +1,6 @@ """Build the sample object.""" -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any from bson import ObjectId # type: ignore @@ -23,11 +23,11 @@ def __init__( self, name: str, experiment: "ExperimentBuilder", - tags: Optional[List[str]] = None, + tags: list[str] | None = None, **metadata, ): self.name = name - self._tasks: List[str] = [] # type: ignore + self._tasks: list[str] = [] # type: ignore self.experiment = experiment self.metadata = metadata self._id = str(ObjectId()) @@ -49,7 +49,7 @@ def add_task( if task_id not in self._tasks: self._tasks.append(task_id) - def to_dict(self) -> Dict[str, Any]: + def to_dict(self) -> dict[str, Any]: """Return Sample as a dictionary. This looks like: diff --git a/alab_management/builders/utils.py b/alab_management/builders/utils.py index 7215ef29..b6f8b8d7 100644 --- a/alab_management/builders/utils.py +++ b/alab_management/builders/utils.py @@ -1,6 +1,6 @@ """This module contains utility functions for the builders module.""" -from typing import TYPE_CHECKING, List, Union +from typing import TYPE_CHECKING from bson import ObjectId # type: ignore @@ -13,7 +13,7 @@ def append_task( task: "BaseTask", - samples: Union[SampleBuilder, List[SampleBuilder]], + samples: SampleBuilder | list[SampleBuilder], ): """ Used to add basetask to a SampleBuilder's tasklist during Experiment construction. diff --git a/alab_management/config.py b/alab_management/config.py index 0177e85f..bd694808 100644 --- a/alab_management/config.py +++ b/alab_management/config.py @@ -27,12 +27,12 @@ import os from pathlib import Path from types import MappingProxyType as FrozenDict -from typing import Any, Dict +from typing import Any import toml -def freeze_config(config_: Dict[str, Any]) -> FrozenDict: +def freeze_config(config_: dict[str, Any]) -> FrozenDict: """ Convert the config dict to frozen config. diff --git a/alab_management/dashboard/routes/experiment.py b/alab_management/dashboard/routes/experiment.py index 469cd8e6..ae94ac97 100644 --- a/alab_management/dashboard/routes/experiment.py +++ b/alab_management/dashboard/routes/experiment.py @@ -1,7 +1,7 @@ """This is a dashboard that displays data from the ALab database.""" from datetime import datetime, timedelta -from typing import Any, Dict, List, Optional +from typing import Any from bson import ObjectId # type: ignore from bson.errors import InvalidId # type: ignore @@ -180,11 +180,11 @@ def query_experiment_results(exp_id: str): def cancel_experiment(exp_id: str): try: exp_id = ObjectId(exp_id) - experiment: Optional[Dict[str, Any]] = experiment_view.get_experiment(exp_id) + experiment: dict[str, Any] | None = experiment_view.get_experiment(exp_id) if experiment is None: return {"status": "error", "reason": "Experiment not found"}, 400 - tasks: List[Dict[str, Any]] = experiment["tasks"] + tasks: list[dict[str, Any]] = experiment["tasks"] # tasks = experiment_view.get_experiment(exp_id)["tasks"] for task in tasks: diff --git a/alab_management/device_manager.py b/alab_management/device_manager.py index bd5ee87a..d45a3064 100644 --- a/alab_management/device_manager.py +++ b/alab_management/device_manager.py @@ -6,11 +6,12 @@ DeviceManager class, which will handle all the request to run certain methods on the real device. """ +from collections.abc import Callable from concurrent.futures import Future from enum import Enum, auto from functools import partial from threading import Thread -from typing import Any, Callable, Dict, NoReturn, Optional, cast +from typing import Any, NoReturn, cast from uuid import uuid4 import dill @@ -171,9 +172,7 @@ def callback_publish(channel, delivery_tag, props, response): channel.basic_ack(delivery_tag=cast(int, delivery_tag)) try: - device_entry: Optional[Dict[str, Any]] = self._device_view.get_device( - device - ) + device_entry: dict[str, Any] | None = self._device_view.get_device(device) # check if the device is currently occupied by this task if self._check_status and ( @@ -217,7 +216,7 @@ def on_message( "kwargs": Dict, } """ - body: Dict[str, Any] = dill.loads(_body) + body: dict[str, Any] = dill.loads(_body) thread = Thread( target=self._execute_command_wrapper, @@ -266,7 +265,7 @@ def __init__(self, task_id: ObjectId, timeout: int = None): # taskid, or can be random? I think this dies with the resourcerequest context manager anyways? self._rpc_reply_queue_name = str(uuid4()) + DEFAULT_CLIENT_QUEUE_SUFFIX self._task_id = task_id - self._waiting: Dict[ObjectId, Future] = {} + self._waiting: dict[ObjectId, Future] = {} self._conn = get_rabbitmq_connection() self._channel = self._conn.channel() @@ -274,7 +273,7 @@ def __init__(self, task_id: ObjectId, timeout: int = None): self._rpc_reply_queue_name, exclusive=False, auto_delete=True ) - self._thread: Optional[Thread] = None + self._thread: Thread | None = None self._channel.basic_consume( queue=self._rpc_reply_queue_name, diff --git a/alab_management/device_view/dbattributes.py b/alab_management/device_view/dbattributes.py index 3c28cfc3..378e7a48 100644 --- a/alab_management/device_view/dbattributes.py +++ b/alab_management/device_view/dbattributes.py @@ -1,4 +1,4 @@ -from typing import Any, Union +from typing import Any from pymongo.collection import Collection # type: ignore @@ -71,7 +71,7 @@ def __init__( device_collection: Collection, device_name: str, attribute_name: str, - default_value: Union[list, None] = None, + default_value: list | None = None, ): self._collection = device_collection self.attribute_name = attribute_name @@ -284,7 +284,7 @@ def __init__( device_collection: Collection, device_name: str, attribute_name: str, - default_value: Union[dict, None] = None, + default_value: dict | None = None, ): self._collection = device_collection self.attribute_name = attribute_name diff --git a/alab_management/device_view/device.py b/alab_management/device_view/device.py index b94a9277..4bdc7bdf 100644 --- a/alab_management/device_view/device.py +++ b/alab_management/device_view/device.py @@ -5,9 +5,10 @@ import threading import time from abc import ABC, abstractmethod +from collections.abc import Callable from queue import Empty, PriorityQueue from traceback import format_exc -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any from unittest.mock import Mock from alab_management.logger import DBLogger @@ -23,7 +24,7 @@ def _UNSPECIFIED(_): def mock( return_constant: Any = _UNSPECIFIED, - object_type: Union[List[Any], Any] = _UNSPECIFIED, + object_type: list[Any] | Any = _UNSPECIFIED, ): """ A decorator used for mocking functions during simulation. @@ -143,7 +144,7 @@ class BaseDevice(ABC): kwargs: keyword arguments that will be passed to the device class """ - def __init__(self, name: str, description: Optional[str] = None, *args, **kwargs): + def __init__(self, name: str, description: str | None = None, *args, **kwargs): """ Initialize a device object, you can set up connection to the device in this method. The device will only be initialized @@ -253,7 +254,7 @@ def disconnect(self): @property @abstractmethod - def sample_positions(self) -> List[SamplePosition]: + def sample_positions(self) -> list[SamplePosition]: """ The sample positions describe the position that can hold a sample. The name of sample position will be the unique identifier of this sample position. It does not store any @@ -298,7 +299,7 @@ def is_running(self) -> bool: # methods to store Device values inside the database. Lists and dictionaries are supported. def list_in_database( - self, name: str, default_value: Optional[Union[list, None]] = None + self, name: str, default_value: list | None | None = None ) -> ListInDatabase: """ Create a list attribute that is stored in the database. @@ -320,7 +321,7 @@ def list_in_database( ) def dict_in_database( - self, name: str, default_value: Optional[Union[dict, None]] = None + self, name: str, default_value: dict | None | None = None ) -> DictInDatabase: """ Create a dict attribute that is stored in the database. @@ -353,7 +354,7 @@ def _apply_default_db_values(self): if any(isinstance(attribute, t) for t in [ListInDatabase, DictInDatabase]): attribute.apply_default_value() - def request_maintenance(self, prompt: str, options: List[Any]): + def request_maintenance(self, prompt: str, options: list[Any]): """ Request maintenance input from the user. This will display a prompt to the user and wait for them to select an option. The selected option will be returned. @@ -365,7 +366,7 @@ def request_maintenance(self, prompt: str, options: List[Any]): return request_maintenance_input(prompt=prompt, options=options) def retrieve_signal( - self, signal_name: str, within: Optional[datetime.timedelta] = None + self, signal_name: str, within: datetime.timedelta | None = None ): """Retrieve a signal from the database. @@ -423,8 +424,8 @@ def __init__(self, device: BaseDevice): self.is_logging = False self.queue: PriorityQueue = PriorityQueue() - self._logging_thread: Optional[threading.Thread] = None - self._start_time: Optional[datetime.datetime] = None + self._logging_thread: threading.Thread | None = None + self._start_time: datetime.datetime | None = None def get_methods_to_log(self): """ @@ -567,7 +568,7 @@ def stop(self): self.is_logging = False self._logging_thread.join() - def retrieve_signal(self, signal_name, within: Optional[datetime.timedelta] = None): + def retrieve_signal(self, signal_name, within: datetime.timedelta | None = None): """Retrieve a signal from the database. Args: @@ -599,7 +600,7 @@ def retrieve_signal(self, signal_name, within: Optional[datetime.timedelta] = No ) -_device_registry: Dict[str, BaseDevice] = {} +_device_registry: dict[str, BaseDevice] = {} def add_device(device: BaseDevice): @@ -609,7 +610,7 @@ def add_device(device: BaseDevice): _device_registry[device.name] = device -def get_all_devices() -> Dict[str, BaseDevice]: +def get_all_devices() -> dict[str, BaseDevice]: """ Get all the device names in the device registry. This is a shallow copy of the registry. diff --git a/alab_management/device_view/device_view.py b/alab_management/device_view/device_view.py index 01fc8366..be414c2b 100644 --- a/alab_management/device_view/device_view.py +++ b/alab_management/device_view/device_view.py @@ -1,8 +1,9 @@ """Wrapper over the ``devices`` collection.""" +from collections.abc import Collection from datetime import datetime from enum import Enum, auto, unique -from typing import Any, Collection, Dict, List, Optional, TypeVar, Union, cast +from typing import Any, TypeVar, cast import pymongo # type: ignore from bson import ObjectId # type: ignore @@ -138,9 +139,9 @@ def add_devices_to_db(self): } ) - def get_all(self) -> List[Dict[str, Any]]: + def get_all(self) -> list[dict[str, Any]]: """Get all the devices in the database, used for dashboard.""" - return cast(List[Dict[str, Any]], self._device_collection.find()) + return cast(list[dict[str, Any]], self._device_collection.find()) def _clean_up_device_collection(self): """Clean up the device collection.""" @@ -149,11 +150,11 @@ def _clean_up_device_collection(self): def request_devices( self, task_id: ObjectId, - device_names_str: Optional[Collection[str]] = None, - device_types_str: Optional[ - Collection[str] - ] = None, # pylint: disable=unsubscriptable-object - ) -> Optional[Dict[str, Dict[str, Union[str, bool]]]]: + device_names_str: Collection[str] | None = None, + device_types_str: ( + Collection[str] | None + ) = None, # pylint: disable=unsubscriptable-object + ) -> dict[str, dict[str, str | bool]] | None: """ Request a list of device, this function will return the name of devices if all the requested device is ready. @@ -182,7 +183,7 @@ def request_devices( "Currently we do not allow duplicated device types in one request." ) - idle_devices: Dict[str, Dict[str, Union[str, bool]]] = {} + idle_devices: dict[str, dict[str, str | bool]] = {} with self._lock(): # pylint: disable=not-callable for device_name in device_names_str: result = self.get_available_devices( @@ -205,8 +206,8 @@ def request_devices( return idle_devices def get_available_devices( - self, device_str: str, type_or_name: str, task_id: Optional[ObjectId] = None - ) -> List[Dict[str, Union[str, bool]]]: + self, device_str: str, type_or_name: str, task_id: ObjectId | None = None + ) -> list[dict[str, str | bool]]: """ Given device type, it will return all the device with this type. @@ -262,7 +263,7 @@ def get_available_devices( for device_entry in self._device_collection.find(request_dict) ] - def get_device(self, device_name: str) -> Dict[str, Any]: + def get_device(self, device_name: str) -> dict[str, Any]: """Get device by device name, if not found, raises ``ValueError``.""" device_entry = self._device_collection.find_one({"name": device_name}) if device_entry is None: @@ -275,7 +276,7 @@ def get_status(self, device_name: str) -> DeviceTaskStatus: return DeviceTaskStatus[device_entry["status"]] - def occupy_device(self, device: Union[BaseDevice, str], task_id: ObjectId): + def occupy_device(self, device: BaseDevice | str, task_id: ObjectId): """Occupy a device with given task id.""" self._update_status( device=device, @@ -284,7 +285,7 @@ def occupy_device(self, device: Union[BaseDevice, str], task_id: ObjectId): task_id=task_id, ) - def get_devices_by_task(self, task_id: Optional[ObjectId]) -> List[BaseDevice]: + def get_devices_by_task(self, task_id: ObjectId | None) -> list[BaseDevice]: """Get devices given a task id (regardless of its status!).""" return [ self._device_list[device["name"]] @@ -333,10 +334,10 @@ def get_samples_on_device(self, device_name: str): def _update_status( self, - device: Union[BaseDevice, str], - required_status: Optional[Union[DeviceTaskStatus, List[DeviceTaskStatus]]], + device: BaseDevice | str, + required_status: DeviceTaskStatus | list[DeviceTaskStatus] | None, target_status: DeviceTaskStatus, - task_id: Optional[ObjectId], + task_id: ObjectId | None, ): """ A method that check and update the status of a device. diff --git a/alab_management/experiment_manager.py b/alab_management/experiment_manager.py index a00992c5..9620172e 100644 --- a/alab_management/experiment_manager.py +++ b/alab_management/experiment_manager.py @@ -7,7 +7,7 @@ """ import time -from typing import Any, Dict, List +from typing import Any from .config import AlabOSConfig from .experiment_view import CompletedExperimentView, ExperimentStatus, ExperimentView @@ -72,9 +72,9 @@ def handle_pending_experiments(self): }, ) - def _handle_pending_experiment(self, experiment: Dict[str, Any]): - samples: List[Dict[str, Any]] = experiment["samples"] - tasks: List[Dict[str, Any]] = experiment["tasks"] + def _handle_pending_experiment(self, experiment: dict[str, Any]): + samples: list[dict[str, Any]] = experiment["samples"] + tasks: list[dict[str, Any]] = experiment["tasks"] # check if there is any cycle in the graph reversed_edges = {i: task["prev_tasks"] for i, task in enumerate(tasks)} diff --git a/alab_management/experiment_view/completed_experiment_view.py b/alab_management/experiment_view/completed_experiment_view.py index b5ee44cb..c59231b3 100644 --- a/alab_management/experiment_view/completed_experiment_view.py +++ b/alab_management/experiment_view/completed_experiment_view.py @@ -1,6 +1,6 @@ """A wrapper over the ``experiment`` class.""" -from typing import Any, Dict, Union +from typing import Any from bson import ObjectId # type: ignore @@ -59,7 +59,7 @@ def save_all(self): except: # noqa: E722 print(f"Error saving experiment {experiment_dict['_id']}") - def exists(self, experiment_id: Union[ObjectId, str]) -> bool: + def exists(self, experiment_id: ObjectId | str) -> bool: """Check if an experiment exists in the completed experiment database. Args: @@ -76,7 +76,7 @@ def exists(self, experiment_id: Union[ObjectId, str]) -> bool: > 0 ) - def get_experiment(self, experiment_id: Union[ObjectId, str]) -> Dict[str, Any]: + def get_experiment(self, experiment_id: ObjectId | str) -> dict[str, Any]: """Get an experiment from the completed experiment collection. Args: diff --git a/alab_management/experiment_view/experiment.py b/alab_management/experiment_view/experiment.py index 0a056d98..d891bfa6 100644 --- a/alab_management/experiment_view/experiment.py +++ b/alab_management/experiment_view/experiment.py @@ -1,6 +1,6 @@ """Define the format of experiment request.""" -from typing import Any, Dict, List, Optional +from typing import Any from bson import BSON, ObjectId # type: ignore from pydantic import ( @@ -12,9 +12,9 @@ class _Sample(BaseModel): name: constr(regex=r"^[^$.]+$") # type: ignore - sample_id: Optional[str] = None - tags: List[str] - metadata: Dict[str, Any] + sample_id: str | None = None + tags: list[str] + metadata: dict[str, Any] @validator("sample_id") def if_provided_must_be_valid_objectid(cls, v): @@ -44,10 +44,10 @@ def must_be_bsonable(cls, v): class _Task(BaseModel): type: str - parameters: Dict[str, Any] - prev_tasks: List[int] - samples: List[str] - task_id: Optional[str] = None + parameters: dict[str, Any] + prev_tasks: list[int] + samples: list[str] + task_id: str | None = None @validator("task_id") def if_provided_must_be_valid_objectid(cls, v): @@ -67,10 +67,10 @@ class InputExperiment(BaseModel): """This is the format that user should follow to write to experiment database.""" name: constr(regex=r"^[^$.]+$") # type: ignore - samples: List[_Sample] - tasks: List[_Task] - tags: List[str] - metadata: Dict[str, Any] + samples: list[_Sample] + tasks: list[_Task] + tags: list[str] + metadata: dict[str, Any] @validator("metadata") def must_be_bsonable(cls, v): diff --git a/alab_management/experiment_view/experiment_view.py b/alab_management/experiment_view/experiment_view.py index 7f311f65..c2f98193 100644 --- a/alab_management/experiment_view/experiment_view.py +++ b/alab_management/experiment_view/experiment_view.py @@ -2,7 +2,7 @@ from datetime import datetime from enum import Enum, auto -from typing import Any, Dict, List, Optional, Union, cast +from typing import Any, cast from bson import ObjectId # type: ignore @@ -84,13 +84,13 @@ def create_experiment(self, experiment: InputExperiment) -> ObjectId: return cast(ObjectId, result.inserted_id) def get_experiments_with_status( - self, status: Union[str, ExperimentStatus] - ) -> List[Dict[str, Any]]: + self, status: str | ExperimentStatus + ) -> list[dict[str, Any]]: """Filter experiments by its status.""" if isinstance(status, str): status = ExperimentStatus[status] return cast( - List[Dict[str, Any]], + list[dict[str, Any]], self._experiment_collection.find( { "status": status.name, @@ -98,7 +98,7 @@ def get_experiments_with_status( ), ) - def get_experiment(self, exp_id: ObjectId) -> Optional[Dict[str, Any]]: + def get_experiment(self, exp_id: ObjectId) -> dict[str, Any] | None: """Get an experiment by its id.""" experiment = self._experiment_collection.find_one({"_id": exp_id}) if experiment is None: @@ -127,7 +127,7 @@ def update_experiment_status(self, exp_id: ObjectId, status: ExperimentStatus): ) def update_sample_task_id( - self, exp_id, sample_ids: List[ObjectId], task_ids: List[ObjectId] + self, exp_id, sample_ids: list[ObjectId], task_ids: list[ObjectId] ): """ At the creation of experiment, the id of samples and tasks has not been assigned. @@ -162,16 +162,14 @@ def update_sample_task_id( }, ) - def get_experiment_by_task_id(self, task_id: ObjectId) -> Optional[Dict[str, Any]]: + def get_experiment_by_task_id(self, task_id: ObjectId) -> dict[str, Any] | None: """Get an experiment that contains a task with the given task_id.""" experiment = self._experiment_collection.find_one({"tasks.task_id": task_id}) if experiment is None: raise ValueError(f"Cannot find experiment containing task_id: {task_id}") return experiment - def get_experiment_by_sample_id( - self, sample_id: ObjectId - ) -> Optional[Dict[str, Any]]: + def get_experiment_by_sample_id(self, sample_id: ObjectId) -> dict[str, Any] | None: """Get an experiment that contains a sample with the given sample_id.""" experiment = self._experiment_collection.find_one( {"samples.sample_id": sample_id} diff --git a/alab_management/lab_view.py b/alab_management/lab_view.py index 90f2e17a..309e5c1e 100644 --- a/alab_management/lab_view.py +++ b/alab_management/lab_view.py @@ -7,7 +7,7 @@ from contextlib import contextmanager from traceback import format_exc -from typing import Any, Dict, List, Optional, Type, Union +from typing import Any from bson import ObjectId from pydantic import root_validator @@ -40,7 +40,7 @@ class ResourcesRequest(BaseModel): :py:class:`SamplePositionRequest ` """ - __root__: Dict[Optional[Type[BaseDevice]], List[SamplePositionRequest]] # type: ignore + __root__: dict[type[BaseDevice] | None, list[SamplePositionRequest]] # type: ignore @root_validator(pre=True, allow_reuse=True) def preprocess(cls, values): @@ -86,11 +86,9 @@ def task_id(self) -> ObjectId: @contextmanager def request_resources( self, - resource_request: Dict[ - Optional[Union[Type[BaseDevice], str]], Dict[str, Union[str, int]] - ], - priority: Optional[int] = None, - timeout: Optional[float] = None, + resource_request: dict[type[BaseDevice] | str | None, dict[str, str | int]], + priority: int | None = None, + timeout: float | None = None, ): """ Request devices and sample positions. This function is a context manager, which should be used in @@ -145,7 +143,7 @@ def _sample_name_to_id(self, sample_name: str) -> ObjectId: f"No sample with name \"{sample_name}\" found for task \"{self.__task_entry['type']}\"" ) - def get_sample(self, sample: Union[ObjectId, str]) -> Sample: + def get_sample(self, sample: ObjectId | str) -> Sample: """ Get a sample by either an ObjectId corresponding to sample_id, or as a string corresponding to the sample's name within the experiment., see also :py:meth:`get_sample @@ -159,7 +157,7 @@ def get_sample(self, sample: Union[ObjectId, str]) -> Sample: raise TypeError("sample must be a sample name (str) or id (ObjectId)") return self._sample_view.get_sample(sample_id=sample_id) - def move_sample(self, sample: Union[ObjectId, str], position: Optional[str]): + def move_sample(self, sample: ObjectId | str, position: str | None): """ Move a sample to a new position. `sample` can be given as either an ObjectId corresponding to sample_id, or as a string corresponding to the sample's name within the experiment. @@ -187,17 +185,15 @@ def move_sample(self, sample: Union[ObjectId, str], position: Optional[str]): sample_id=sample_entry.sample_id, position=position ) - def get_locked_sample_positions(self) -> List[str]: + def get_locked_sample_positions(self) -> list[str]: """Get a list of sample positions that are occupied by this task.""" return self._sample_view.get_sample_positions_by_task(task_id=self._task_id) - def get_sample_position_parent_device(self, position: str) -> Optional[str]: + def get_sample_position_parent_device(self, position: str) -> str | None: """Get the name of the device that owns the sample position.""" return self._sample_view.get_sample_position_parent_device(position=position) - def update_sample_metadata( - self, sample: Union[ObjectId, str], metadata: Dict[str, Any] - ): + def update_sample_metadata(self, sample: ObjectId | str, metadata: dict[str, Any]): """ Update the metadata of a sample. `sample` can be given as either an ObjectId corresponding to sample_id, or as a string corresponding to the sample's name within the experiment. @@ -218,7 +214,7 @@ def update_sample_metadata( ) def run_subtask( - self, task: Type[BaseTask], samples: List[Union[ObjectId, str]], **kwargs + self, task: type[BaseTask], samples: list[ObjectId | str], **kwargs ): """Run a task as a subtask within the task. basically fills in task_id and lab_view for you. this command blocks until the subtask is completed. @@ -318,7 +314,7 @@ def run_subtask( ) return result - def request_user_input(self, prompt: str, options: List[str]) -> str: + def request_user_input(self, prompt: str, options: list[str]) -> str: """Request user input from the user. This function will block until the user inputs something. Returns the value returned by the user. """ diff --git a/alab_management/logger.py b/alab_management/logger.py index 937d4d21..5db83cc3 100644 --- a/alab_management/logger.py +++ b/alab_management/logger.py @@ -1,8 +1,9 @@ """Logger module takes charge of recording information, warnings and errors during executing tasks.""" +from collections.abc import Iterable from datetime import datetime, timedelta from enum import Enum, auto, unique -from typing import Any, Dict, Iterable, Optional, Union, cast +from typing import Any, cast from bson import ObjectId @@ -35,14 +36,14 @@ class LoggingLevel(Enum): class DBLogger: """A custom logger that wrote data to database, where we predefined some log pattern.""" - def __init__(self, task_id: Optional[ObjectId]): + def __init__(self, task_id: ObjectId | None): self.task_id = task_id self._logging_collection = get_collection("logs") def log( self, - level: Union[str, int, LoggingLevel], - log_data: Dict[str, Any], + level: str | int | LoggingLevel, + log_data: dict[str, Any], logging_type: LoggingType = LoggingType.OTHER, ) -> ObjectId: """ @@ -70,7 +71,7 @@ def log( return cast(ObjectId, result.inserted_id) - def log_amount(self, log_data: Dict[str, Any]): + def log_amount(self, log_data: dict[str, Any]): """Log the amount of samples and chemicals (e.g. weight).""" return self.log( level=LoggingLevel.INFO, @@ -78,7 +79,7 @@ def log_amount(self, log_data: Dict[str, Any]): logging_type=LoggingType.SAMPLE_AMOUNT, ) - def log_characterization_result(self, log_data: Dict[str, Any]): + def log_characterization_result(self, log_data: dict[str, Any]): """Log the characterization result (e.g. XRD pattern).""" return self.log( level=LoggingLevel.INFO, @@ -98,17 +99,15 @@ def log_device_signal(self, device_name: str, signal_name: str, signal_value: An logging_type=LoggingType.DEVICE_SIGNAL, ) - def system_log( - self, level: Union[str, int, LoggingLevel], log_data: Dict[str, Any] - ): + def system_log(self, level: str | int | LoggingLevel, log_data: dict[str, Any]): """Log that comes from the workflow system.""" return self.log( level=level, log_data=log_data, logging_type=LoggingType.SYSTEM_LOG ) def filter_log( - self, level: Union[str, int, LoggingLevel], within: timedelta - ) -> Iterable[Dict[str, Any]]: + self, level: str | int | LoggingLevel, within: timedelta + ) -> Iterable[dict[str, Any]]: """Find log within a range of time (1h/1d or else) higher than certain level.""" if isinstance(level, str): level = cast(int, LoggingLevel[level].value) @@ -121,7 +120,7 @@ def filter_log( def get_latest_device_signal( self, device_name: str, signal_name: str - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Get the last device signal log. Args: @@ -163,7 +162,7 @@ def get_latest_device_signal( def filter_device_signal( self, device_name: str, signal_name: str, within: timedelta - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Find device signal log within a range of time (1h/1d or else). Args: diff --git a/alab_management/sample_view/completed_sample_view.py b/alab_management/sample_view/completed_sample_view.py index 18fdce30..151fd0c4 100644 --- a/alab_management/sample_view/completed_sample_view.py +++ b/alab_management/sample_view/completed_sample_view.py @@ -3,8 +3,6 @@ saving samples to the completed database. """ -from typing import Union - from bson import ObjectId # type: ignore from alab_management.utils.data_objects import get_collection, get_completed_collection @@ -42,7 +40,7 @@ def save_sample(self, sample_id: ObjectId): else: self._completed_sample_collection.insert_one(sample_dict) - def exists(self, sample_id: Union[ObjectId, str]) -> bool: + def exists(self, sample_id: ObjectId | str) -> bool: """Check if a sample exists in the database. Args: diff --git a/alab_management/sample_view/sample.py b/alab_management/sample_view/sample.py index 6db3adc8..bbe498b2 100644 --- a/alab_management/sample_view/sample.py +++ b/alab_management/sample_view/sample.py @@ -1,7 +1,7 @@ """The definition of the Sample and SamplePosition classes.""" from dataclasses import dataclass, field -from typing import Any, ClassVar, Dict, List, Optional +from typing import Any, ClassVar from bson import ObjectId # type: ignore @@ -19,11 +19,11 @@ class Sample: """ sample_id: ObjectId - task_id: Optional[ObjectId] + task_id: ObjectId | None name: str - position: Optional[str] - metadata: Dict[str, Any] = field(default_factory=dict) - tags: List[str] = field(default_factory=list) + position: str | None + metadata: dict[str, Any] = field(default_factory=dict) + tags: list[str] = field(default_factory=list) @dataclass(frozen=True) @@ -54,7 +54,7 @@ def __post_init__(self): ) -_standalone_sample_position_registry: Dict[str, SamplePosition] = {} +_standalone_sample_position_registry: dict[str, SamplePosition] = {} def add_standalone_sample_position(position: SamplePosition): @@ -68,6 +68,6 @@ def add_standalone_sample_position(position: SamplePosition): _standalone_sample_position_registry[position.name] = position -def get_all_standalone_sample_positions() -> Dict[str, SamplePosition]: +def get_all_standalone_sample_positions() -> dict[str, SamplePosition]: """Get all the device names in the device registry.""" return _standalone_sample_position_registry.copy() diff --git a/alab_management/sample_view/sample_view.py b/alab_management/sample_view/sample_view.py index 9e029862..e76451c9 100644 --- a/alab_management/sample_view/sample_view.py +++ b/alab_management/sample_view/sample_view.py @@ -3,7 +3,7 @@ import re from datetime import datetime from enum import Enum, auto -from typing import Any, Dict, List, Optional, Tuple, Union, cast +from typing import Any, cast import pymongo # type: ignore from bson import ObjectId # type: ignore @@ -36,7 +36,7 @@ def from_str(cls, sample_position_prefix: str) -> "SamplePositionRequest": return cls(prefix=sample_position_prefix) @classmethod - def from_py_type(cls, sample_position: Union[str, Dict[str, Any]]): + def from_py_type(cls, sample_position: str | dict[str, Any]): """Create a ``SamplePositionRequest`` from a string or a dict.""" if isinstance(sample_position, str): return cls.from_str(sample_position_prefix=sample_position) @@ -75,8 +75,8 @@ def __init__(self): def add_sample_positions_to_db( self, - sample_positions: List[SamplePosition], - parent_device_name: Optional[str] = None, + sample_positions: list[SamplePosition], + parent_device_name: str | None = None, ): """ Insert sample positions info to db, which includes position name and description. @@ -126,8 +126,8 @@ def clean_up_sample_position_collection(self): def request_sample_positions( self, task_id: ObjectId, - sample_positions: List[Union[SamplePositionRequest, str, Dict[str, Any]]], - ) -> Optional[Dict[str, List[Dict[str, Any]]]]: + sample_positions: list[SamplePositionRequest | str | dict[str, Any]], + ) -> dict[str, list[dict[str, Any]]] | None: """ Request a list of sample positions, this function will return until all the sample positions are available. @@ -137,7 +137,7 @@ def request_sample_positions( The sample position name is actually the prefix of a sample position, which we will try to match all the sample positions will the name """ - sample_positions_request: List[SamplePositionRequest] = [ + sample_positions_request: list[SamplePositionRequest] = [ ( SamplePositionRequest.from_py_type(sample_position) if not isinstance(sample_position, SamplePositionRequest) @@ -163,7 +163,7 @@ def request_sample_positions( ) with self._lock(): # pylint: disable=not-callable - available_positions: Dict[str, List[Dict[str, Union[str, bool]]]] = {} + available_positions: dict[str, list[dict[str, str | bool]]] = {} for sample_position in sample_positions_request: result = self.get_available_sample_position( task_id, position_prefix=sample_position.prefix @@ -176,7 +176,7 @@ def request_sample_positions( )[: sample_position.number] return available_positions - def get_sample_position(self, position: str) -> Optional[Dict[str, Any]]: + def get_sample_position(self, position: str) -> dict[str, Any] | None: """ Get the sample position entry in the database. @@ -186,7 +186,7 @@ def get_sample_position(self, position: str) -> Optional[Dict[str, Any]]: def get_sample_position_status( self, position: str - ) -> Tuple[SamplePositionStatus, Optional[ObjectId]]: + ) -> tuple[SamplePositionStatus, ObjectId | None]: """ Get the status of a sample position. @@ -216,7 +216,7 @@ def get_sample_position_status( return SamplePositionStatus.EMPTY, None - def get_sample_position_parent_device(self, position: str) -> Optional[str]: + def get_sample_position_parent_device(self, position: str) -> str | None: """ Get the parent device of a sample position. @@ -247,7 +247,7 @@ def is_unoccupied_position(self, position: str) -> bool: def get_available_sample_position( self, task_id: ObjectId, position_prefix: str - ) -> List[Dict[str, Union[str, bool]]]: + ) -> list[dict[str, str | bool]]: """ Check if the position is occupied. @@ -325,7 +325,7 @@ def release_sample_position(self, position: str): }, ) - def get_sample_positions_by_task(self, task_id: Optional[ObjectId]) -> List[str]: + def get_sample_positions_by_task(self, task_id: ObjectId | None) -> list[str]: """Get the list of sample positions that is locked by a task (given task id).""" return [ sample_position["name"] @@ -341,10 +341,10 @@ def get_sample_positions_by_task(self, task_id: Optional[ObjectId]) -> List[str] def create_sample( self, name: str, - position: Optional[str] = None, - sample_id: Optional[ObjectId] = None, - tags: Optional[List[str]] = None, - metadata: Optional[Dict[str, Any]] = None, + position: str | None = None, + sample_id: ObjectId | None = None, + tags: list[str] | None = None, + metadata: dict[str, Any] | None = None, ) -> ObjectId: """ Create a sample and return its uid in the database. @@ -408,7 +408,7 @@ def get_sample(self, sample_id: ObjectId) -> Sample: tags=result.get("tags", []), ) - def update_sample_task_id(self, sample_id: ObjectId, task_id: Optional[ObjectId]): + def update_sample_task_id(self, sample_id: ObjectId, task_id: ObjectId | None): """Update the task id for a sample.""" result = self._sample_collection.find_one({"_id": sample_id}) if result is None: @@ -424,7 +424,7 @@ def update_sample_task_id(self, sample_id: ObjectId, task_id: Optional[ObjectId] }, ) - def update_sample_metadata(self, sample_id: ObjectId, metadata: Dict[str, Any]): + def update_sample_metadata(self, sample_id: ObjectId, metadata: dict[str, Any]): """Update the metadata for a sample. This adds new metadata or updates existing metadata.""" result = self._sample_collection.find_one({"_id": sample_id}) if result is None: @@ -438,7 +438,7 @@ def update_sample_metadata(self, sample_id: ObjectId, metadata: Dict[str, Any]): {"$set": update_dict}, ) - def move_sample(self, sample_id: ObjectId, position: Optional[str]): + def move_sample(self, sample_id: ObjectId, position: str | None): """Update the sample with new position.""" result = self._sample_collection.find_one({"_id": sample_id}) if result is None: @@ -462,7 +462,7 @@ def move_sample(self, sample_id: ObjectId, position: Optional[str]): }, ) - def exists(self, sample_id: Union[ObjectId, str]) -> bool: + def exists(self, sample_id: ObjectId | str) -> bool: """Check if a sample exists in the database. Args: diff --git a/alab_management/task_manager/resource_requester.py b/alab_management/task_manager/resource_requester.py index d2b36771..b79a59e2 100644 --- a/alab_management/task_manager/resource_requester.py +++ b/alab_management/task_manager/resource_requester.py @@ -8,7 +8,7 @@ from datetime import datetime from threading import Thread from traceback import print_exc -from typing import Any, Dict, List, Optional, Type, Union, cast +from typing import Any, cast import dill from bson import ObjectId @@ -22,9 +22,9 @@ from .enums import _EXTRA_REQUEST, RequestStatus -_SampleRequestDict = Dict[str, int] -_ResourceRequestDict = Dict[ - Optional[Union[Type[BaseDevice], str]], _SampleRequestDict +_SampleRequestDict = dict[str, int] +_ResourceRequestDict = dict[ + type[BaseDevice] | str | None, _SampleRequestDict ] # the raw request sent by task process @@ -53,8 +53,8 @@ class ResourcesRequest(BaseModel): :py:class:`SamplePositionRequest ` """ - __root__: List[ - Dict[str, Union[List[Dict[str, Union[str, int]]], Dict[str, str]]] + __root__: list[ + dict[str, list[dict[str, str | int]] | dict[str, str]] ] # type: ignore @root_validator(pre=True, allow_reuse=True) @@ -126,9 +126,9 @@ def __init__( task_id: ObjectId, ): self._request_collection = get_collection("requests") - self._waiting: Dict[ObjectId, Dict[str, Any]] = {} + self._waiting: dict[ObjectId, dict[str, Any]] = {} self.task_id = task_id - self.priority: Union[int, TaskPriority] = ( + self.priority: int | TaskPriority = ( TaskPriority.NORMAL ) # will usually be overwritten by BaseTask instantiation. @@ -150,9 +150,9 @@ def __close__(self): def request_resources( self, resource_request: _ResourceRequestDict, - timeout: Optional[float] = None, - priority: Optional[Union[TaskPriority, int]] = None, - ) -> Dict[str, Any]: + timeout: float | None = None, + priority: TaskPriority | int | None = None, + ) -> dict[str, Any]: """ Request lab resources. @@ -293,13 +293,13 @@ def _handle_fulfilled_request(self, request_id: ObjectId): if entry["status"] != RequestStatus.FULFILLED.name: # type: ignore return - assigned_devices: Dict[str, Dict[str, Union[str, bool]]] = entry["assigned_devices"] # type: ignore - assigned_sample_positions: Dict[str, List[Dict[str, Any]]] = entry["assigned_sample_positions"] # type: ignore + assigned_devices: dict[str, dict[str, str | bool]] = entry["assigned_devices"] # type: ignore + assigned_sample_positions: dict[str, list[dict[str, Any]]] = entry["assigned_sample_positions"] # type: ignore - request: Dict[str, Any] = self._waiting.pop(request_id) + request: dict[str, Any] = self._waiting.pop(request_id) f: Future = request["f"] - device_str_to_request: Dict[str, Union[Type[BaseDevice], str, None]] = request[ + device_str_to_request: dict[str, type[BaseDevice] | str | None] = request[ "device_str_to_request" ] @@ -326,18 +326,18 @@ def _handle_error_request(self, request_id: ObjectId): return error: Exception = dill.loads(entry["error"]) # type: ignore - request: Dict[str, Any] = self._waiting.pop(request_id) + request: dict[str, Any] = self._waiting.pop(request_id) f: Future = request["f"] f.set_exception(error) @staticmethod def _post_process_requested_resource( - devices: Dict[Type[BaseDevice], str], - sample_positions: Dict[str, List[str]], - resource_request: Dict[str, List[Dict[str, Union[int, str]]]], + devices: dict[type[BaseDevice], str], + sample_positions: dict[str, list[str]], + resource_request: dict[str, list[dict[str, int | str]]], ): - processed_sample_positions: Dict[ - Optional[Type[BaseDevice]], Dict[str, List[str]] + processed_sample_positions: dict[ + type[BaseDevice] | None, dict[str, list[str]] ] = {} for device_request, sample_position_dict in resource_request.items(): diff --git a/alab_management/task_manager/task_manager.py b/alab_management/task_manager/task_manager.py index ff7ebd0c..0f4ab276 100644 --- a/alab_management/task_manager/task_manager.py +++ b/alab_management/task_manager/task_manager.py @@ -8,7 +8,7 @@ from functools import partial from math import inf from threading import Thread -from typing import Any, Dict, List, Type, cast +from typing import Any, cast import dill import networkx as nx @@ -36,7 +36,7 @@ ) -def parse_reroute_tasks() -> Dict[str, Type[BaseTask]]: +def parse_reroute_tasks() -> dict[str, type[BaseTask]]: """ Takes the reroute task registry and expands the supported sample positions (which is given in format similar to resource requests) to the individual sample positions. @@ -56,7 +56,7 @@ def parse_reroute_tasks() -> Dict[str, Type[BaseTask]]: load_definition() - routes: Dict[str, BaseTask] = {} # sample_position: Task + routes: dict[str, BaseTask] = {} # sample_position: Task sample_view = SampleView() for reroute in _reroute_task_registry: @@ -294,7 +294,7 @@ def handle_request_cycles(self): thread.daemon = False thread.start() - def _handle_requested_resources(self, request_entry: Dict[str, Any]): + def _handle_requested_resources(self, request_entry: dict[str, Any]): try: resource_request = request_entry["request"] task_id = request_entry["task_id"] @@ -396,14 +396,14 @@ def _handle_requested_resources(self, request_entry: Dict[str, Any]): sample_positions=sample_positions, task_id=task_id ) - def _occupy_devices(self, devices: Dict[str, Dict[str, Any]], task_id: ObjectId): + def _occupy_devices(self, devices: dict[str, dict[str, Any]], task_id: ObjectId): for device in devices.values(): self.device_view.occupy_device( device=cast(str, device["name"]), task_id=task_id ) def _occupy_sample_positions( - self, sample_positions: Dict[str, List[Dict[str, Any]]], task_id: ObjectId + self, sample_positions: dict[str, list[dict[str, Any]]], task_id: ObjectId ): for sample_positions_ in sample_positions.values(): for sample_position_ in sample_positions_: @@ -411,13 +411,13 @@ def _occupy_sample_positions( task_id, cast(str, sample_position_["name"]) ) - def _release_devices(self, devices: Dict[str, Dict[str, Any]]): + def _release_devices(self, devices: dict[str, dict[str, Any]]): for device in devices.values(): if device["need_release"]: self.device_view.release_device(device["name"]) def _release_sample_positions( - self, sample_positions: Dict[str, List[Dict[str, Any]]] + self, sample_positions: dict[str, list[dict[str, Any]]] ): for sample_positions_ in sample_positions.values(): for sample_position in sample_positions_: @@ -509,7 +509,7 @@ def _check_for_request_cycle(self): def _reroute_to_fix_request_cycle( self, task_id: ObjectId, - sample_positions: List[str], + sample_positions: list[str], ): from alab_management.lab_view import LabView diff --git a/alab_management/task_view/completed_task_view.py b/alab_management/task_view/completed_task_view.py index 98b95365..f5dbb7ec 100644 --- a/alab_management/task_view/completed_task_view.py +++ b/alab_management/task_view/completed_task_view.py @@ -3,8 +3,6 @@ completed task database. """ -from typing import Union - from bson import ObjectId from alab_management.utils.data_objects import get_collection, get_completed_collection @@ -41,7 +39,7 @@ def save_task(self, task_id: ObjectId): else: self._completed_task_collection.insert_one(task_dict) - def exists(self, task_id: Union[ObjectId, str]) -> bool: + def exists(self, task_id: ObjectId | str) -> bool: """ Check if a task exists in the database. diff --git a/alab_management/task_view/task.py b/alab_management/task_view/task.py index 6d6774b2..64eec233 100644 --- a/alab_management/task_view/task.py +++ b/alab_management/task_view/task.py @@ -3,7 +3,7 @@ import inspect from abc import ABC, abstractmethod from inspect import getfullargspec -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union +from typing import TYPE_CHECKING, Any, Optional from bson.objectid import ObjectId @@ -32,7 +32,7 @@ def to_json(self): } @classmethod - def from_json(cls, json: Dict[str, Any]): + def from_json(cls, json: dict[str, Any]): """Create a ResultPointer from a JSON-serializable dictionary.""" if json["type"] != "ResultPointer": raise ValueError("JSON does not encode a ResultPointer!") @@ -48,10 +48,10 @@ class BaseTask(ABC): def __init__( self, - samples: Optional[List[Union[str, ObjectId]]] = None, - task_id: Optional[ObjectId] = None, + samples: list[str | ObjectId] | None = None, + task_id: ObjectId | None = None, lab_view: Optional["LabView"] = None, - priority: Optional[Union[TaskPriority, int]] = TaskPriority.NORMAL, + priority: TaskPriority | int | None = TaskPriority.NORMAL, simulation: bool = True, *args, **kwargs, @@ -108,7 +108,7 @@ def is_simulation(self) -> bool: return self.__simulation @property - def samples(self) -> List[str]: + def samples(self) -> list[str]: """Returns the list of samples associated with this task.""" return self.__samples @@ -121,7 +121,7 @@ def priority(self) -> int: @property # @abstractmethod - def result_specification(self) -> Dict[str, Any]: + def result_specification(self) -> dict[str, Any]: """Returns a dictionary describing the results to be generated by this task. Raises @@ -191,7 +191,7 @@ def export_result(self, key: str) -> dict: def import_result( self, - pointer: Union[ResultPointer, Dict[str, Any]], + pointer: ResultPointer | dict[str, Any], allow_explicit_value: bool = False, ) -> Any: """ @@ -232,7 +232,7 @@ def import_result( return reference_task["result"][pointer.key] @priority.setter - def priority(self, value: Union[int, TaskPriority]): + def priority(self, value: int | TaskPriority): if value < 0: raise ValueError("Priority should be a positive integer") if not self.__simulation: @@ -316,8 +316,8 @@ def run(self): def run_subtask( self, - task: Type["BaseTask"], - samples: Optional[Union[List[str], str]] = None, + task: type["BaseTask"], + samples: list[str] | str | None = None, **kwargs, ): """Run a subtask of this current task. Returns the result, if any, of the subtask.""" @@ -328,7 +328,7 @@ def run_subtask( def add_to( self, - samples: Union[SampleBuilder, List[SampleBuilder]], + samples: SampleBuilder | list[SampleBuilder], ): """Used to add basetask to a SampleBuilder's tasklist during Experiment construction. @@ -355,31 +355,29 @@ def add_to( sample.add_task(task_id=task_id) -_task_registry: Dict[str, Type[BaseTask]] = {} +_task_registry: dict[str, type[BaseTask]] = {} -SUPPORTED_SAMPLE_POSITIONS_TYPE = Dict[ - Union[Type["BaseDevice"], str, None], Union[str, List[str]] -] -_reroute_task_registry: List[ - Dict[str, Union[Type[BaseTask], SUPPORTED_SAMPLE_POSITIONS_TYPE]] +SUPPORTED_SAMPLE_POSITIONS_TYPE = dict[type["BaseDevice"] | str | None, str | list[str]] +_reroute_task_registry: list[ + dict[str, type[BaseTask] | SUPPORTED_SAMPLE_POSITIONS_TYPE] ] = [] -def add_task(task: Type[BaseTask]): +def add_task(task: type[BaseTask]): """Register a task.""" if task.__name__ in _task_registry: raise KeyError(f"Duplicated operation name {task.__name__}") _task_registry[task.__name__] = task -def get_all_tasks() -> Dict[str, Type[BaseTask]]: +def get_all_tasks() -> dict[str, type[BaseTask]]: """Get all the tasks in the registry.""" return _task_registry.copy() def add_reroute_task( supported_sample_positions: SUPPORTED_SAMPLE_POSITIONS_TYPE, - task: Type[BaseTask], + task: type[BaseTask], **kwargs, ): """Register a reroute task.""" diff --git a/alab_management/task_view/task_view.py b/alab_management/task_view/task_view.py index 4ea99d66..94393840 100644 --- a/alab_management/task_view/task_view.py +++ b/alab_management/task_view/task_view.py @@ -4,7 +4,7 @@ """ from datetime import datetime -from typing import Any, Dict, List, Optional, Type, Union, cast +from typing import Any, cast from bson import ObjectId @@ -23,16 +23,16 @@ class TaskView: def __init__(self): self._task_collection = get_collection("tasks") self._lock = get_lock("tasks") - self._tasks_definition: Dict[str, Type[BaseTask]] = get_all_tasks() + self._tasks_definition: dict[str, type[BaseTask]] = get_all_tasks() def create_task( self, task_type: str, - samples: List[ObjectId], - parameters: Dict[str, Any], - prev_tasks: Optional[Union[ObjectId, List[ObjectId]]] = None, - next_tasks: Optional[Union[ObjectId, List[ObjectId]]] = None, - task_id: Optional[ObjectId] = None, + samples: list[ObjectId], + parameters: dict[str, Any], + prev_tasks: ObjectId | list[ObjectId] | None = None, + next_tasks: ObjectId | list[ObjectId] | None = None, + task_id: ObjectId | None = None, ) -> ObjectId: """ Insert a task into the task collection. @@ -81,7 +81,7 @@ def create_task( return cast(ObjectId, result.inserted_id) def create_subtask( - self, task_id, subtask_type, samples: List[str], parameters: dict + self, task_id, subtask_type, samples: list[str], parameters: dict ): """Create a subtask entry for a task.""" task = self.get_task(task_id=task_id) @@ -111,7 +111,7 @@ def create_subtask( ) return subtask_id - def get_task(self, task_id: ObjectId, encode: bool = False) -> Dict[str, Any]: + def get_task(self, task_id: ObjectId, encode: bool = False) -> dict[str, Any]: """ Get a task by its task id, which will return all the info stored in the database. @@ -137,7 +137,7 @@ def get_task(self, task_id: ObjectId, encode: bool = False) -> Dict[str, Any]: result = self.encode_task(result) return result - def get_task_with_sample(self, sample_id: ObjectId) -> Optional[Dict[str, Any]]: + def get_task_with_sample(self, sample_id: ObjectId) -> dict[str, Any] | None: """Get a task that contains the sample with the provided id.""" result = self._task_collection.find({"samples.sample_id": sample_id}) if result is None: @@ -257,7 +257,7 @@ def update_subtask_status( ) def update_result( - self, task_id: ObjectId, name: Optional[str] = None, value: Any = None + self, task_id: ObjectId, name: str | None = None, value: Any = None ): """ Update result to completed job. @@ -335,7 +335,7 @@ def try_to_mark_task_ready(self, task_id: ObjectId): ): self.update_status(task_id, TaskStatus.READY) - def get_ready_tasks(self) -> List[Dict[str, Any]]: + def get_ready_tasks(self) -> list[dict[str, Any]]: """ Return a list of ready tasks. @@ -346,7 +346,7 @@ def get_ready_tasks(self) -> List[Dict[str, Any]]: """ return self.get_tasks_by_status(status=TaskStatus.READY) - def get_tasks_by_status(self, status: TaskStatus) -> List[Dict[str, Any]]: + def get_tasks_by_status(self, status: TaskStatus) -> list[dict[str, Any]]: """ Return a list of tasks with given status. @@ -357,17 +357,17 @@ def get_tasks_by_status(self, status: TaskStatus) -> List[Dict[str, Any]]: """ result = self._task_collection.find({"status": status.name}) - tasks: List[Dict[str, Any]] = [] + tasks: list[dict[str, Any]] = [] for task_entry in result: tasks.append(self.encode_task(task_entry)) return tasks - def encode_task(self, task_entry: Dict[str, Any]) -> Dict[str, Any]: + def encode_task(self, task_entry: dict[str, Any]) -> dict[str, Any]: """ Rename _id to task_id Translate task's type into corresponding python class. """ - operation_type: Type[BaseTask] = self._tasks_definition[task_entry["type"]] + operation_type: type[BaseTask] = self._tasks_definition[task_entry["type"]] task_entry["task_id"] = task_entry.pop( "_id" ) # change the key name of `_id` to `task_id` @@ -379,8 +379,8 @@ def encode_task(self, task_entry: Dict[str, Any]) -> Dict[str, Any]: def update_task_dependency( self, task_id: ObjectId, - prev_tasks: Optional[Union[ObjectId, List[ObjectId]]] = None, - next_tasks: Optional[Union[ObjectId, List[ObjectId]]] = None, + prev_tasks: ObjectId | list[ObjectId] | None = None, + next_tasks: ObjectId | list[ObjectId] | None = None, ): """ Add prev tasks and next tasks to one task entry, @@ -474,6 +474,6 @@ def mark_task_as_cancelling(self, task_id: ObjectId): status=TaskStatus.CANCELLING, ) - def exists(self, task_id: Union[ObjectId, str]) -> bool: + def exists(self, task_id: ObjectId | str) -> bool: """Check if a task id exists.""" return self._task_collection.count_documents({"_id": ObjectId(task_id)}) > 0 diff --git a/alab_management/user_input.py b/alab_management/user_input.py index 9db8cacc..6e359b5e 100644 --- a/alab_management/user_input.py +++ b/alab_management/user_input.py @@ -1,6 +1,6 @@ import time from enum import Enum -from typing import Any, Dict, List, Optional, Union, cast +from typing import Any, cast from bson import ObjectId @@ -35,8 +35,8 @@ def __init__(self): def insert_request( self, prompt: str, - options: List[str], - task_id: Optional[ObjectId] = None, + options: list[str], + task_id: ObjectId | None = None, maintenance: bool = False, category: str = "Unknown Category", ) -> ObjectId: @@ -76,7 +76,7 @@ def insert_request( self._alarm.alert(f"User input requested: {prompt}", category) return request_id - def get_request(self, request_id: ObjectId) -> Dict[str, Any]: + def get_request(self, request_id: ObjectId) -> dict[str, Any]: """ Get a request. @@ -86,7 +86,7 @@ def get_request(self, request_id: ObjectId) -> Dict[str, Any]: if request is None: raise ValueError(f"User input request id {request_id} does not exist!") - return cast(Dict[str, Any], request) + return cast(dict[str, Any], request) def update_request_status(self, request_id: ObjectId, response: str, note: str): """Update the status of a request.""" @@ -128,15 +128,15 @@ def get_all_pending_requests(self) -> list: Returns a list of pending requests. """ return cast( - List[Dict[str, Any]], + list[dict[str, Any]], self._input_collection.find({"status": UserRequestStatus.PENDING.value}), ) def request_user_input( - task_id: Union[ObjectId, None], + task_id: ObjectId | None, prompt: str, - options: List[str], + options: list[str], maintenance: bool = False, category: str = "Unknown Category", ) -> str: @@ -161,7 +161,7 @@ def request_user_input( return user_input_view.retrieve_user_input(request_id=request_id) -def request_maintenance_input(prompt: str, options: List[str]): +def request_maintenance_input(prompt: str, options: list[str]): """ Request user input through the dashboard. Blocks until response is given. diff --git a/alab_management/utils/data_objects.py b/alab_management/utils/data_objects.py index 87a9ef9c..d4f8e386 100644 --- a/alab_management/utils/data_objects.py +++ b/alab_management/utils/data_objects.py @@ -5,7 +5,6 @@ from abc import ABC, abstractmethod from datetime import datetime from enum import Enum -from typing import Optional import numpy as np import pika @@ -40,9 +39,9 @@ def get_lock(cls, name: str) -> MongoLock: class _GetMongoCollection(_BaseGetMongoCollection): - client: Optional[pymongo.MongoClient] = None - db: Optional[database.Database] = None - db_lock: Optional[MongoLock] = None + client: pymongo.MongoClient | None = None + db: database.Database | None = None + db_lock: MongoLock | None = None @classmethod def init(cls): @@ -62,9 +61,9 @@ def init(cls): class _GetCompletedMongoCollection(_BaseGetMongoCollection): - client: Optional[pymongo.MongoClient] = None - db: Optional[database.Database] = None - db_lock: Optional[MongoLock] = None + client: pymongo.MongoClient | None = None + db: database.Database | None = None + db_lock: MongoLock | None = None @classmethod def init(cls): diff --git a/alab_management/utils/db_lock.py b/alab_management/utils/db_lock.py index bac01050..3756598d 100644 --- a/alab_management/utils/db_lock.py +++ b/alab_management/utils/db_lock.py @@ -2,7 +2,6 @@ import time from contextlib import contextmanager -from typing import Optional from pymongo.collection import Collection from pymongo.errors import DuplicateKeyError @@ -30,12 +29,12 @@ def name(self) -> str: return self._name @contextmanager - def __call__(self, timeout: Optional[float] = None): + def __call__(self, timeout: float | None = None): """Acquire the lock and release it after the context is finished.""" yield self.acquire(timeout=timeout) self.release() - def acquire(self, timeout: Optional[float] = None): + def acquire(self, timeout: float | None = None): """Acquire the lock.""" start_time = time.time() while timeout is None or time.time() - start_time <= timeout: diff --git a/alab_management/utils/graph_ops.py b/alab_management/utils/graph_ops.py index aa81b27c..c02d990c 100644 --- a/alab_management/utils/graph_ops.py +++ b/alab_management/utils/graph_ops.py @@ -1,12 +1,12 @@ """This file contains a custom graph class and functions for checking if there are cycles in a graph.""" -from typing import Any, Dict, List +from typing import Any class Graph: """Use adjacent table to store a graph.""" - def __init__(self, vertices: List[Any], edges: Dict[int, List[int]]): + def __init__(self, vertices: list[Any], edges: dict[int, list[int]]): """Note that the all the keys and values are the index of vertices.""" if len(vertices) != len(edges): raise ValueError( @@ -45,7 +45,7 @@ def _is_cyclic(v): for vertex in range(len(self.vertices)) ) - def get_parents(self, v: Any) -> List[Any]: + def get_parents(self, v: Any) -> list[Any]: """Provide the value of vertex, return the value of its parents vertices.""" index = self.vertices.index(v) return [ @@ -55,7 +55,7 @@ def get_parents(self, v: Any) -> List[Any]: if child == index ] - def get_children(self, v: Any) -> List[Any]: + def get_children(self, v: Any) -> list[Any]: """Provide the index of vertex, return the value of its children vertices.""" index = self.vertices.index(v) return [self.vertices[i] for i in self.edges[index]] diff --git a/alab_management/utils/module_ops.py b/alab_management/utils/module_ops.py index abadedaa..3a4b1645 100644 --- a/alab_management/utils/module_ops.py +++ b/alab_management/utils/module_ops.py @@ -5,12 +5,9 @@ import sys from copy import copy from pathlib import Path -from typing import Optional, Union -def import_module_from_path( - path: Union[str, Path], parent_package: Optional[str] = None -): +def import_module_from_path(path: str | Path, parent_package: str | None = None): """Import a module by its path.""" if not isinstance(path, Path): path = Path(path) diff --git a/examples/fake_lab/tasks/ending.py b/examples/fake_lab/tasks/ending.py index ffa3e124..7c77599b 100644 --- a/examples/fake_lab/tasks/ending.py +++ b/examples/fake_lab/tasks/ending.py @@ -1,12 +1,10 @@ -from typing import List - from bson import ObjectId from alab_management.task_view.task import BaseTask class Ending(BaseTask): - def __init__(self, samples: List[ObjectId], *args, **kwargs): + def __init__(self, samples: list[ObjectId], *args, **kwargs): super().__init__(samples=samples, *args, **kwargs) self.sample = samples[0] diff --git a/examples/fake_lab/tasks/heating.py b/examples/fake_lab/tasks/heating.py index c43a7efb..66f5a541 100644 --- a/examples/fake_lab/tasks/heating.py +++ b/examples/fake_lab/tasks/heating.py @@ -1,5 +1,4 @@ import time -from typing import List, Tuple from bson import ObjectId @@ -12,8 +11,8 @@ class Heating(BaseTask): def __init__( self, - samples: List[ObjectId], - setpoints: List[Tuple[float, float]], + samples: list[ObjectId], + setpoints: list[tuple[float, float]], *args, **kwargs, ): diff --git a/examples/fake_lab/tasks/moving.py b/examples/fake_lab/tasks/moving.py index 7bdbe0f8..25d3ed28 100644 --- a/examples/fake_lab/tasks/moving.py +++ b/examples/fake_lab/tasks/moving.py @@ -1,4 +1,4 @@ -from typing import List, cast +from typing import cast from bson import ObjectId @@ -7,7 +7,7 @@ class Moving(BaseTask): - def __init__(self, samples: List[ObjectId], dest: str, *args, **kwargs): + def __init__(self, samples: list[ObjectId], dest: str, *args, **kwargs): super().__init__(samples=samples, *args, **kwargs) self.sample = samples[0] self.dest = dest diff --git a/examples/fake_lab/tasks/starting.py b/examples/fake_lab/tasks/starting.py index 3c0c5030..9ba265af 100644 --- a/examples/fake_lab/tasks/starting.py +++ b/examples/fake_lab/tasks/starting.py @@ -1,5 +1,4 @@ import time -from typing import List from bson import ObjectId @@ -7,7 +6,7 @@ class Starting(BaseTask): - def __init__(self, samples: List[ObjectId], dest: str, *args, **kwargs): + def __init__(self, samples: list[ObjectId], dest: str, *args, **kwargs): super().__init__(samples=samples, *args, **kwargs) self.sample = samples[0] self.dest = dest diff --git a/pyproject.toml b/pyproject.toml index 83cd19af..4576b756 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -132,7 +132,7 @@ max-doc-length = 400 # jobflow uses 88. Ensure to change this back to 88 in the select = "C, E, F, W, B" extend-ignore = "E203, W503, E501, F401, RST21, D102" ignore = "E203, W503, E501, F401, RST21, D102" -min-python-version = "3.8.0" +min-python-version = "3.10.0" docstring-convention = "numpy" rst-roles = "class, func, ref, obj" @@ -163,7 +163,7 @@ exclude_lines = [ ] [tool.ruff] -target-version = "py38" +target-version = "py310" ignore-init-module-imports = true select = [ "B", # flake8-bugbear diff --git a/tests/fake_lab/tasks/ending.py b/tests/fake_lab/tasks/ending.py index ffa3e124..7c77599b 100644 --- a/tests/fake_lab/tasks/ending.py +++ b/tests/fake_lab/tasks/ending.py @@ -1,12 +1,10 @@ -from typing import List - from bson import ObjectId from alab_management.task_view.task import BaseTask class Ending(BaseTask): - def __init__(self, samples: List[ObjectId], *args, **kwargs): + def __init__(self, samples: list[ObjectId], *args, **kwargs): super().__init__(samples=samples, *args, **kwargs) self.sample = samples[0] diff --git a/tests/fake_lab/tasks/heating.py b/tests/fake_lab/tasks/heating.py index dd08699e..e80d1935 100644 --- a/tests/fake_lab/tasks/heating.py +++ b/tests/fake_lab/tasks/heating.py @@ -1,5 +1,4 @@ import time -from typing import List, Tuple from bson import ObjectId @@ -12,8 +11,8 @@ class Heating(BaseTask): def __init__( self, - samples: List[ObjectId], - setpoints: List[Tuple[float, float]], + samples: list[ObjectId], + setpoints: list[tuple[float, float]], *args, **kwargs, ): diff --git a/tests/fake_lab/tasks/moving.py b/tests/fake_lab/tasks/moving.py index 3e1752f2..9fbd0eb3 100644 --- a/tests/fake_lab/tasks/moving.py +++ b/tests/fake_lab/tasks/moving.py @@ -1,4 +1,4 @@ -from typing import List, cast +from typing import cast from bson import ObjectId @@ -8,7 +8,7 @@ class Moving(BaseTask): - def __init__(self, samples: List[ObjectId], dest: str, *args, **kwargs): + def __init__(self, samples: list[ObjectId], dest: str, *args, **kwargs): super().__init__(samples=samples, *args, **kwargs) self.sample = samples[0] self.dest = dest diff --git a/tests/fake_lab/tasks/starting.py b/tests/fake_lab/tasks/starting.py index 3c0c5030..9ba265af 100644 --- a/tests/fake_lab/tasks/starting.py +++ b/tests/fake_lab/tasks/starting.py @@ -1,5 +1,4 @@ import time -from typing import List from bson import ObjectId @@ -7,7 +6,7 @@ class Starting(BaseTask): - def __init__(self, samples: List[ObjectId], dest: str, *args, **kwargs): + def __init__(self, samples: list[ObjectId], dest: str, *args, **kwargs): super().__init__(samples=samples, *args, **kwargs) self.sample = samples[0] self.dest = dest From 6269612b21ff5a78b796971ac57da2df01c6faf9 Mon Sep 17 00:00:00 2001 From: Bernardus Rendy <37468936+bernardusrendy@users.noreply.github.com> Date: Sun, 17 Mar 2024 14:02:23 -0700 Subject: [PATCH 8/9] Update cli.py to support dev mode dev mode: pip install -e . --- alab_management/scripts/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/alab_management/scripts/cli.py b/alab_management/scripts/cli.py index ea8f0971..7a487f14 100644 --- a/alab_management/scripts/cli.py +++ b/alab_management/scripts/cli.py @@ -2,7 +2,7 @@ import click -from alab_management import __version__ +from alab_management.__init__ import __version__ from alab_management.config import AlabOSConfig from .cleanup_lab import cleanup_lab From ed1b67b4f0da2c7de757acc52d979ecb92c0f439 Mon Sep 17 00:00:00 2001 From: Bernardus Rendy <37468936+bernardusrendy@users.noreply.github.com> Date: Sun, 17 Mar 2024 14:43:08 -0700 Subject: [PATCH 9/9] Imports to allow editable mode --- alab_management/_default/devices/default_device.py | 3 ++- alab_management/builders/utils.py | 2 +- alab_management/task_manager/task_manager.py | 3 +-- docs/source/conf.py | 2 +- examples/fake_lab/__init__.py | 9 +++------ examples/fake_lab/devices/furnace.py | 3 ++- examples/fake_lab/devices/robot_arm.py | 3 ++- examples/fake_lab/tasks/heating.py | 2 +- tests/fake_lab/__init__.py | 9 +++------ tests/fake_lab/devices/furnace.py | 3 ++- tests/fake_lab/devices/robot_arm.py | 3 ++- tests/fake_lab/tasks/heating.py | 2 +- 12 files changed, 21 insertions(+), 23 deletions(-) diff --git a/alab_management/_default/devices/default_device.py b/alab_management/_default/devices/default_device.py index 73472e0b..614cccf4 100644 --- a/alab_management/_default/devices/default_device.py +++ b/alab_management/_default/devices/default_device.py @@ -1,6 +1,7 @@ from typing import ClassVar -from alab_management import BaseDevice, SamplePosition +from alab_management.device_view import BaseDevice +from alab_management.sample_view import SamplePosition class DefaultDevice(BaseDevice): diff --git a/alab_management/builders/utils.py b/alab_management/builders/utils.py index b6f8b8d7..821bad51 100644 --- a/alab_management/builders/utils.py +++ b/alab_management/builders/utils.py @@ -8,7 +8,7 @@ from .samplebuilder import SampleBuilder if TYPE_CHECKING: - from alab_management import BaseTask + from alab_management.task_view import BaseTask def append_task( diff --git a/alab_management/task_manager/task_manager.py b/alab_management/task_manager/task_manager.py index 0f4ab276..832b27dc 100644 --- a/alab_management/task_manager/task_manager.py +++ b/alab_management/task_manager/task_manager.py @@ -15,8 +15,7 @@ from bson import ObjectId from dramatiq_abort import abort -from alab_management import BaseDevice -from alab_management.device_view import get_all_devices +from alab_management.device_view import BaseDevice, get_all_devices from alab_management.device_view.device_view import DeviceView from alab_management.lab_view import LabView from alab_management.logger import DBLogger, LoggingLevel diff --git a/docs/source/conf.py b/docs/source/conf.py index a7d3c3c4..1d4fa4cb 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -30,7 +30,7 @@ author = "Alab Project Team" # The full version, including alpha/beta/rc tags -from alab_management import __version__ # noqa +from alab_management.__init__ import __version__ # noqa release = __version__ diff --git a/examples/fake_lab/__init__.py b/examples/fake_lab/__init__.py index 3edb5611..15431605 100644 --- a/examples/fake_lab/__init__.py +++ b/examples/fake_lab/__init__.py @@ -1,9 +1,6 @@ -from alab_management import ( - SamplePosition, - add_device, - add_standalone_sample_position, - add_task, -) +from alab_management.device_view import add_device +from alab_management.sample_view import SamplePosition, add_standalone_sample_position +from alab_management.task_view import add_task from .devices.furnace import Furnace from .devices.robot_arm import RobotArm diff --git a/examples/fake_lab/devices/furnace.py b/examples/fake_lab/devices/furnace.py index 2323c4cc..16e2f8cf 100644 --- a/examples/fake_lab/devices/furnace.py +++ b/examples/fake_lab/devices/furnace.py @@ -1,7 +1,8 @@ from threading import Timer from typing import ClassVar -from alab_management import BaseDevice, SamplePosition +from alab_management.device_view import BaseDevice +from alab_management.sample_view import SamplePosition class Furnace(BaseDevice): diff --git a/examples/fake_lab/devices/robot_arm.py b/examples/fake_lab/devices/robot_arm.py index 977d30fb..48e5cf03 100644 --- a/examples/fake_lab/devices/robot_arm.py +++ b/examples/fake_lab/devices/robot_arm.py @@ -1,6 +1,7 @@ from typing import ClassVar -from alab_management import BaseDevice, SamplePosition +from alab_management.device_view import BaseDevice +from alab_management.sample_view import SamplePosition class RobotArm(BaseDevice): diff --git a/examples/fake_lab/tasks/heating.py b/examples/fake_lab/tasks/heating.py index 66f5a541..c953b668 100644 --- a/examples/fake_lab/tasks/heating.py +++ b/examples/fake_lab/tasks/heating.py @@ -2,7 +2,7 @@ from bson import ObjectId -from alab_management import BaseTask +from alab_management.task_view import BaseTask from fake_lab.devices.furnace import Furnace from .moving import Moving diff --git a/tests/fake_lab/__init__.py b/tests/fake_lab/__init__.py index 3edb5611..15431605 100644 --- a/tests/fake_lab/__init__.py +++ b/tests/fake_lab/__init__.py @@ -1,9 +1,6 @@ -from alab_management import ( - SamplePosition, - add_device, - add_standalone_sample_position, - add_task, -) +from alab_management.device_view import add_device +from alab_management.sample_view import SamplePosition, add_standalone_sample_position +from alab_management.task_view import add_task from .devices.furnace import Furnace from .devices.robot_arm import RobotArm diff --git a/tests/fake_lab/devices/furnace.py b/tests/fake_lab/devices/furnace.py index 2323c4cc..16e2f8cf 100644 --- a/tests/fake_lab/devices/furnace.py +++ b/tests/fake_lab/devices/furnace.py @@ -1,7 +1,8 @@ from threading import Timer from typing import ClassVar -from alab_management import BaseDevice, SamplePosition +from alab_management.device_view import BaseDevice +from alab_management.sample_view import SamplePosition class Furnace(BaseDevice): diff --git a/tests/fake_lab/devices/robot_arm.py b/tests/fake_lab/devices/robot_arm.py index 977d30fb..48e5cf03 100644 --- a/tests/fake_lab/devices/robot_arm.py +++ b/tests/fake_lab/devices/robot_arm.py @@ -1,6 +1,7 @@ from typing import ClassVar -from alab_management import BaseDevice, SamplePosition +from alab_management.device_view import BaseDevice +from alab_management.sample_view import SamplePosition class RobotArm(BaseDevice): diff --git a/tests/fake_lab/tasks/heating.py b/tests/fake_lab/tasks/heating.py index e80d1935..d30fbbb7 100644 --- a/tests/fake_lab/tasks/heating.py +++ b/tests/fake_lab/tasks/heating.py @@ -2,7 +2,7 @@ from bson import ObjectId -from alab_management import BaseTask +from alab_management.task_view import BaseTask from ..devices.furnace import Furnace # noqa from .moving import Moving