From 1f831684432d0134f500749561b50ce8373d3a87 Mon Sep 17 00:00:00 2001 From: Patrick Latimer Date: Wed, 5 Jun 2024 12:08:36 -0700 Subject: [PATCH] style: fix lint complaints and add docstrings --- .gitignore | 1 + src/aind_slims_api/__init__.py | 5 +- src/aind_slims_api/configuration.py | 7 ++- src/aind_slims_api/core.py | 90 ++++++++++++++++++++--------- src/aind_slims_api/mouse.py | 11 ++-- src/aind_slims_api/user.py | 11 +++- src/aind_slims_api/waterlog.py | 70 ++++++++++++++++------ tests/test_slimsclient.py | 15 ++--- tests/test_waterlog.py | 18 +++--- 9 files changed, 155 insertions(+), 73 deletions(-) diff --git a/.gitignore b/.gitignore index 06a56dd..df28b4e 100644 --- a/.gitignore +++ b/.gitignore @@ -109,6 +109,7 @@ venv/ ENV/ env.bak/ venv.bak/ +.conda # Spyder project settings .spyderproject diff --git a/src/aind_slims_api/__init__.py b/src/aind_slims_api/__init__.py index 28c0f78..246bdf2 100644 --- a/src/aind_slims_api/__init__.py +++ b/src/aind_slims_api/__init__.py @@ -1,9 +1,10 @@ """Init package""" +from .core import SlimsClient # noqa + + __version__ = "0.0.0" from .configuration import AindSlimsApiSettings config = AindSlimsApiSettings() - -from .core import SlimsClient diff --git a/src/aind_slims_api/configuration.py b/src/aind_slims_api/configuration.py index f1a6a5b..a017ff9 100644 --- a/src/aind_slims_api/configuration.py +++ b/src/aind_slims_api/configuration.py @@ -1,13 +1,14 @@ -from typing import Optional +""" Library Configuration model """ -from pydantic import HttpUrl, SecretStr +from pydantic import SecretStr from pydantic_settings import BaseSettings class AindSlimsApiSettings(BaseSettings): """Settings for SLIMS Client - Per pydantic-settings docs https://docs.pydantic.dev/latest/concepts/pydantic_settings/ + Per pydantic-settings docs + https://docs.pydantic.dev/latest/concepts/pydantic_settings/ Loads slims credentials from environment variables if present""" slims_url: str = "https://aind-test.us.slims.agilent.com/slimsrest/" diff --git a/src/aind_slims_api/core.py b/src/aind_slims_api/core.py index b5c2756..17d5691 100644 --- a/src/aind_slims_api/core.py +++ b/src/aind_slims_api/core.py @@ -1,28 +1,24 @@ -from datetime import timezone as tz -import datetime +"""Contents: + +Utilities for creating pydantic models for SLIMS data: + SlimsBaseModel - to be subclassed for SLIMS pydantic models + UnitSpec - To be included in a type annotation of a Quantity field + +SlimsClient - Basic wrapper around slims-python-api client with convenience + methods and integration with SlimsBaseModel subtypes +""" + from datetime import datetime -from enum import Enum from functools import lru_cache -import os -import json -from dataclasses import dataclass -from zoneinfo import ZoneInfo from pydantic import ( BaseModel, - BeforeValidator, - ValidationError, ValidationInfo, - conlist, field_serializer, field_validator, - model_validator, ) from pydantic.fields import FieldInfo -import requests -from requests.auth import HTTPBasicAuth import logging -from typing import Annotated, Any, Self, Union, Literal, Optional -import math +from typing import Literal, Optional from slims.slims import Slims, _SlimsApiException from slims.internal import ( @@ -49,10 +45,13 @@ class UnitSpec: + """Used in type annotation metadata to specify units""" + units: list[str] preferred_unit: str = None def __init__(self, *args, preferred_unit=None): + """Set list of acceptable units from args, and preferred_unit""" self.units = args if len(self.units) == 0: raise ValueError("One or more units must be specified") @@ -60,7 +59,8 @@ def __init__(self, *args, preferred_unit=None): self.preferred_unit = self.units[0] -def find_unit_spec(field: FieldInfo) -> UnitSpec | None: +def _find_unit_spec(field: FieldInfo) -> UnitSpec | None: + """Given a Pydantic FieldInfo, find the UnitSpec in its metadata""" metadata = field.metadata for m in metadata: if isinstance(m, UnitSpec): @@ -90,29 +90,36 @@ class MyModel(SlimsBaseModel): _slims_table: SLIMSTABLES @field_validator("*", mode="before") - def validate(cls, value, info: ValidationInfo): + def _validate(cls, value, info: ValidationInfo): + """Validates a field, accounts for Quantities""" if isinstance(value, SlimsColumn): if value.datatype == "QUANTITY": - unit_spec = find_unit_spec(cls.model_fields[info.field_name]) + unit_spec = _find_unit_spec(cls.model_fields[info.field_name]) if unit_spec is None: - msg = f'Quantity field "{info.field_name}" must be annotated with a UnitSpec' + msg = ( + f'Quantity field "{info.field_name}"' + "must be annotated with a UnitSpec" + ) raise TypeError(msg) if value.unit not in unit_spec.units: - msg = f'Unexpected unit "{value.unit}" for field {info.field_name}, Expected {unit_spec.units}' + msg = ( + f'Unexpected unit "{value.unit}" for field ' + f"{info.field_name}, Expected {unit_spec.units}" + ) raise ValueError(msg) return value.value else: return value @field_serializer("*") - def serialize(self, field, info): - unit_spec = find_unit_spec(self.model_fields[info.field_name]) + def _serialize(self, field, info): + """Serialize a field, accounts for Quantities and datetime""" + unit_spec = _find_unit_spec(self.model_fields[info.field_name]) if unit_spec and field is not None: quantity = { "amount": field, "unit_display": unit_spec.preferred_unit, } - # quantity["unit_pk"] = 6 if unit_spec.preferred_unit == "g" else 15 return quantity elif isinstance(field, datetime): return int(field.timestamp() * 10**3) @@ -125,7 +132,10 @@ def serialize(self, field, info): class SlimsClient: + """Wrapper around slims-python-api client with convenience methods""" + def __init__(self, url=None, username=None, password=None): + """Create object and try to connect to database""" self.url = url or config.slims_url self.db: Slims = None @@ -136,6 +146,7 @@ def __init__(self, url=None, username=None, password=None): ) def connect(self, url: str, username: str, password: str): + """Connect to the database""" self.db = Slims( "slims", url, @@ -186,11 +197,9 @@ def fetch( return records - def fetch_unit(self, unit_name: str) -> SlimsRecord: - return self.fetch("Unit", unit_name=unit_name)[0] - @lru_cache(maxsize=None) def fetch_pk(self, table: SLIMSTABLES, *args, **kwargs) -> int | None: + """SlimsClient.fetch but returns the pk of the first returned record""" records = self.fetch(table, *args, **kwargs) if len(records) > 0: return records[0].pk() @@ -198,14 +207,17 @@ def fetch_pk(self, table: SLIMSTABLES, *args, **kwargs) -> int | None: return None def fetch_user(self, user_name: str): + """Fetches a user by username""" return self.fetch("User", user_userName=user_name) def add(self, table: SLIMSTABLES, data: dict): + """Add a SLIMS record to a given SLIMS table""" record = self.db.add(table, data) logger.info(f"SLIMS Add: {table}/{record.pk()}") return record def update(self, table: SLIMSTABLES, pk: int, data: dict): + """Update a SLIMS record""" record = self.db.fetch_by_pk(table, pk) if record is None: raise ValueError('No data in SLIMS "{table}" table for pk "{pk}"') @@ -214,11 +226,24 @@ def update(self, table: SLIMSTABLES, pk: int, data: dict): return new_record def rest_link(self, table: SLIMSTABLES, **kwargs): + """Construct a url link to a SLIMS table with arbitrary filters""" base_url = f"{self.url}/rest/{table}" queries = [f"?{k}={v}" for k, v in kwargs.items()] return base_url + "".join(queries) - def add_model(self, model: SlimsBaseModel, *args, **kwargs): + def add_model( + self, model: SlimsBaseModel, *args, **kwargs + ) -> SlimsBaseModel: + """Given a SlimsBaseModel object, add it to SLIMS + Args + model (SlimsBaseModel): object to add + *args (str): fields to include in the serialization + **kwargs: passed to model.model_dump() + + Returns + An instance of the same type of model, with data from + the resulting SLIMS record + """ fields_to_include = set(args) or None fields_to_exclude = set(kwargs.get("exclude", [])) fields_to_exclude.add("pk") @@ -234,6 +259,17 @@ def add_model(self, model: SlimsBaseModel, *args, **kwargs): return type(model).model_validate(rtn) def update_model(self, model: SlimsBaseModel, *args, **kwargs): + """Given a SlimsBaseModel object, update its (existing) SLIMS record + + Args + model (SlimsBaseModel): object to update + *args (str): fields to include in the serialization + **kwargs: passed to model.model_dump() + + Returns + An instance of the same type of model, with data from + the resulting SLIMS record + """ fields_to_include = set(args) or None rtn = self.update( model._slims_table, diff --git a/src/aind_slims_api/mouse.py b/src/aind_slims_api/mouse.py index 7c29971..983dcb4 100644 --- a/src/aind_slims_api/mouse.py +++ b/src/aind_slims_api/mouse.py @@ -1,7 +1,7 @@ -import logging -from typing import Annotated, Optional +"""Contains a model for the mouse content, and a method for fetching it""" -from slims.slims import Record +import logging +from typing import Annotated from pydantic import Field, ValidationError @@ -11,6 +11,8 @@ class SlimsMouseContent(SlimsBaseModel): + """Model for an instance of the Mouse ContentType""" + baseline_weight_g: Annotated[float | None, UnitSpec("g")] = Field( ..., alias="cntn_cf_baselineWeight" ) @@ -55,7 +57,8 @@ def fetch_mouse_content( mouse_details = mice[0] if len(mice) > 1: logger.warning( - f"Warning, Multiple mice in SLIMS with barcode {mouse_name}, using pk={mouse_details.cntn_pk}" + f"Warning, Multiple mice in SLIMS with barcode " + f"{mouse_name}, using pk={mouse_details.cntn_pk}" ) else: logger.warning("Warning, Mouse not in SLIMS") diff --git a/src/aind_slims_api/user.py b/src/aind_slims_api/user.py index 9f88566..0715599 100644 --- a/src/aind_slims_api/user.py +++ b/src/aind_slims_api/user.py @@ -1,15 +1,19 @@ +"""Contains a model for a user, and a method for fetching it""" + import logging -from typing import Annotated, Optional +from typing import Optional from pydantic import Field, ValidationError -from .core import SlimsBaseModel, SlimsClient, UnitSpec +from .core import SlimsBaseModel, SlimsClient logger = logging.getLogger() # TODO: Tighten this up once users are more commonly used class SlimsUser(SlimsBaseModel): + """Model for user information in SLIMS""" + username: str = Field(..., alias="user_userName") first_name: Optional[str] = Field("", alias="user_firstName") last_name: Optional[str] = Field("", alias="user_lastName") @@ -34,7 +38,8 @@ def fetch_user( user_details = users[0] if len(users) > 1: logger.warning( - f"Warning, Multiple users in SLIMS with username {users}, using pk={user_details.pk}" + f"Warning, Multiple users in SLIMS with " + f"username {users}, using pk={user_details.pk}" ) else: logger.warning("Warning, User not in SLIMS") diff --git a/src/aind_slims_api/waterlog.py b/src/aind_slims_api/waterlog.py index 9650bcb..296ef4a 100644 --- a/src/aind_slims_api/waterlog.py +++ b/src/aind_slims_api/waterlog.py @@ -1,17 +1,15 @@ +"""Contains models for Waterlog Results and Water Restriction Events and +methods for fetching each of those. + +Also contains a Mouse class for keeping track of relevant data for a given +mouse and fetching/updating/posting data to SLIMS""" + from datetime import datetime -from functools import partial import logging import os -from typing import Annotated, Optional, Literal +from typing import Annotated, Optional -from pydantic import ( - BaseModel, - BeforeValidator, - Field, - ValidationError, - confloat, -) -from slims.slims import _SlimsApiException +from pydantic import Field, ValidationError import aind_slims_api from aind_slims_api.user import SlimsUser @@ -23,6 +21,8 @@ class SLIMSWaterlogResult(SlimsBaseModel): + """Model for a SLIMS Waterlog Result, the daily water/weight records""" + date: datetime = Field(datetime.now(), alias="rslt_cf_datePerformed") operator: Optional[str] = Field(None, alias="rslt_cf_waterlogOperator") weight_g: Annotated[float | None, UnitSpec("g")] = Field( @@ -54,6 +54,8 @@ class SLIMSWaterlogResult(SlimsBaseModel): class SlimsWaterRestrictionEvent(SlimsBaseModel): + """Model for a Water Restriction Event""" + start_date: datetime = Field(datetime.now(), alias="cnvn_cf_startDate") end_date: Optional[datetime] = Field(None, alias="cnvn_cf_endDate") assigned_by: str = Field(..., alias="cnvn_cf_assignedBy") @@ -74,7 +76,7 @@ def fetch_mouse_waterlog_results( client: SlimsClient, mouse: SlimsMouseContent, ) -> list[SLIMSWaterlogResult]: - """Fetch "Waterlog" Results from SLIMS and update the mouse accordingly + """Fetch "Waterlog" Results from SLIMS Args: client (SlimsClient): SLIMS client object @@ -113,7 +115,7 @@ def fetch_water_restriction_events( mouse (SlimsMouseContent): Mouse data object Returns: - list[SlimsWaterRestrictionEvent]: SLIMS records of water restriction events + list[SlimsWaterRestrictionEvent]: SLIMS records of events """ slims_records: list[SlimsWaterRestrictionEvent] = client.fetch( @@ -135,7 +137,10 @@ def fetch_water_restriction_events( class Mouse: + """Class for tracking mouse water/weight data and syncing with SLIMS""" + def __init__(self, mouse_name: str, user: SlimsUser, slims_client=None): + """Fetch data from Slims for mouse with barcode={mouse_name}""" self.client = slims_client or SlimsClient() self.mouse_name = mouse_name self.user = user @@ -145,10 +150,15 @@ def __init__(self, mouse_name: str, user: SlimsUser, slims_client=None): self.restriction: SlimsWaterRestrictionEvent = None self.all_restrictions: list[SlimsWaterRestrictionEvent] = [] + self.link_mouse: str = None + self.link_restrictions: str = None + self.link_wl_records: str = None + self._fetch_data() - self.get_pks() + self._fetch_pks() def _fetch_data(self): + """Fetches mouse/waterlog/restriction data from SLIMS""" self.mouse = fetch_mouse_content(self.client, self.mouse_name) self.waterlog_results = fetch_mouse_waterlog_results( @@ -163,13 +173,15 @@ def _fetch_data(self): event_active = latest_restriction.end_date is None if not (self.mouse.water_restricted == event_active): logger.warning( - f"Warning, inconsistent water restricted data in SLIMS, MID, {self.mouse.barcode}" + f"Warning, inconsistent water restricted data in SLIMS, " + f"MID, {self.mouse.barcode}" ) self.restriction = latest_restriction if event_active else None - self.make_links() + self._make_links() - def make_links(self): + def _make_links(self): + """Constructs useful links to SLIMS tables""" self.link_mouse = self.client.rest_link( "Content", cntn_cf_labtracksId=self.mouse_name ) @@ -180,7 +192,8 @@ def make_links(self): "ContentEvent", cnvn_fk_content=self.mouse_name ) - def get_pks(self): + def _fetch_pks(self): + """Fetches useful SLIMS pks""" self.wrest_pk = self.client.fetch_pk( "ContentEventType", cnvt_uniqueIdentifier="cnvt_water_restriction" ) @@ -197,6 +210,8 @@ def add_waterlog_record( total_water: float = None, comments: str = None, ): + """Creates and adds a new waterlog weight/water record to SLIMS, and + updates self.waterlog_results accordingly""" if total_water is None: total_water = water_earned + water_supplement_delivered @@ -217,13 +232,20 @@ def add_waterlog_record( logger.info("Added SLIMS Waterlog record") def post_baseline_weight(self, new_baseline_weight: float): + """Update the baseline weight in SLIMS and self.mouse""" self.mouse.baseline_weight_g = new_baseline_weight self.mouse = self.client.update_model(self.mouse, "baseline_weight") logger.info( - f"Updated mouse {self.mouse_name} baseline weight to {new_baseline_weight}" + f"Updated mouse {self.mouse_name} " + f"baseline weight to {new_baseline_weight}" ) def switch_to_water_restricted(self, target_weight_fraction: float): + """If the mouse is on ad-lib water, + - Create a water restriction event starting today + - Update the mouse's water_restricted field + - Update SLIMS with the above + - Update local data accordingly""" if self.mouse.water_restricted: logger.info("Mouse is already water restricted") return @@ -242,6 +264,11 @@ def switch_to_water_restricted(self, target_weight_fraction: float): logger.info(f"Switched mouse {self.mouse_name} to Water Restricted") def switch_to_adlib_water(self): + """If the mouse is water restricted, + - Set the end date of the active restriction event to today + - Update the mouse's water_restricted field + - Update SLIMS with the above + - Update local data accordingly""" if not self.mouse.water_restricted: logger.info("Mouse is already on ad-lib water") return @@ -257,10 +284,15 @@ def switch_to_adlib_water(self): logger.info(f"Switched mouse {self.mouse_name} to Adlib Water") def update_target_weight_fraction(self, new_twf: float): + """Update the target weight fraction of the active restriction""" + if not self.mouse.water_restricted: + logger.info("Mouse is not water restricted") + return self.restriction.target_weight_fraction = new_twf self.restriction = self.client.update_model( self.restriction, "target_weight_fraction" ) logger.info( - f"Updated mouse {self.mouse_name} target weight fraction to {new_twf}" + f"Updated mouse {self.mouse_name} " + f"target weight fraction to {new_twf}" ) diff --git a/tests/test_slimsclient.py b/tests/test_slimsclient.py index d7c4387..b9b942c 100644 --- a/tests/test_slimsclient.py +++ b/tests/test_slimsclient.py @@ -1,18 +1,19 @@ +""" Tests the generic SlimsClient object""" + import unittest -from aind_slims_api.configuration import AindSlimsApiSettings -from aind_slims_api.core import SlimsClient -from aind_slims_api.waterlog import WaterlogSlimsClient +# from aind_slims_api.configuration import AindSlimsApiSettings +# from aind_slims_api.core import SlimsClient class TestSlimsClient(unittest.TestCase): """Example Test Class""" - def test_config(self): - config = AindSlimsApiSettings() + # def test_config(self): + # config = AindSlimsApiSettings() - def test_slims_client(self): - client = SlimsClient() + # def test_slims_client(self): + # client = SlimsClient() # @pytest.mark.parametrize("mouse_name", ["614173"]) # def test_waterlog(mouse_name): diff --git a/tests/test_waterlog.py b/tests/test_waterlog.py index dd90371..c168813 100644 --- a/tests/test_waterlog.py +++ b/tests/test_waterlog.py @@ -1,7 +1,7 @@ +""" Tests waterlog models, methods, and Mouse class""" + import unittest -from aind_slims_api.configuration import AindSlimsApiSettings -import datetime from aind_slims_api import SlimsClient from aind_slims_api.waterlog import Mouse from aind_slims_api.user import SlimsUser @@ -11,6 +11,7 @@ class TestWaterlog(unittest.TestCase): """Example Test Class""" def setup(self): + """set client, user, and test_mouse_id""" self.client = SlimsClient() # self.user = fetch_user(client, "SIPE") self.user = SlimsUser( @@ -21,6 +22,7 @@ def setup(self): self.test_mouse_id = "614173" def test_main_sequence_using_mouse_object(self): + """Runs through waterlog methods""" self.setup() mouse = Mouse( @@ -30,7 +32,7 @@ def test_main_sequence_using_mouse_object(self): if mouse.mouse.water_restricted: mouse.switch_to_adlib_water() - ## Make some waterlog entries + # Make some waterlog entries for i in range(3): mouse.add_waterlog_record( weight=20 + i, @@ -40,13 +42,13 @@ def test_main_sequence_using_mouse_object(self): comments=f"Test {i}", ) - ## Post baseline weight + # Post baseline weight mouse.post_baseline_weight(21) - ## Water restrict + # Water restrict mouse.switch_to_water_restricted(target_weight_fraction=0.86) - ## Add another waterlog result + # Add another waterlog result mouse.add_waterlog_record( weight=20, water_earned=0.5, @@ -55,10 +57,10 @@ def test_main_sequence_using_mouse_object(self): comments=f"Test {i+1}", ) - ## Update target weight fraction + # Update target weight fraction mouse.update_target_weight_fraction(0.9) - ## Switch to adlib + # Switch to adlib mouse.switch_to_adlib_water()