Skip to content

Commit

Permalink
refactor: Updated the get_model_series() function to make it simple…
Browse files Browse the repository at this point in the history 
…r for future maintenance and additions. (#145)
  • Loading branch information
@nfelt14
nfelt14 authored Feb 14, 2024
1 parent 70d5da6 commit 5369d67
Show file tree
Hide file tree
Showing 5 changed files with 155 additions and 97 deletions.
4 changes: 4 additions & 0 deletions CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,10 @@ ______________________________________________________________________

Things to be included in the next release go here.

### Changed

- Updated the `get_model_series()` function to use a regex mapping instead of complicated logic to reduce maintenance costs.

______________________________________________________________________

## v1.2.0 (2024-02-09)
Expand Down
17 changes: 10 additions & 7 deletions docs/contributing/add_new_driver.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,23 +25,26 @@ This guide will walk through the steps needed to add a new device driver.
- See other `__init__.py` files for examples
04. Update the `SupportedModels` enum exposed in
`tm_devices/helpers/__init__.py`
05. Update the `DEVICE_DRIVER_MODEL_MAPPING` lookup inside
05. Update the `_SUPPORTED_MODEL_REGEX_MAPPING` regex constant inside
`tm_devices/helpers/functions.py` to include a mapping of the new driver name (model series)
to a regex string matching the appropriate model strings
06. Update the `DEVICE_DRIVER_MODEL_MAPPING` lookup inside
`tm_devices/drivers/__init__.py`
06. Update the `__all__` variable inside `tm_devices/drivers/__init__.py` to
07. Update the `__all__` variable inside `tm_devices/drivers/__init__.py` to
include the new device driver
07. Update the appropriate Type Alias in `device_manager.py` (search for "Type
08. Update the appropriate Type Alias in `device_manager.py` (search for "Type
Aliases")
08. If the device supports VISA USBTMC communication, update the
09. If the device supports VISA USBTMC communication, update the
`USB_MODEL_ID_LOOKUP` lookup exposed in `tm_devices/helpers/__init__.py`
09. Update the Supported Devices section in `README.rst` to include the new model
10. Update unit tests (and simulated device files)
10. Update the Supported Devices section in `README.rst` to include the new model
11. Update unit tests (and simulated device files)
1. Add a new simulated device driver in the correct folder within
`tests/sim_devices`
2. Update `tests/sim_devices/devices.yaml` with a new resource for the new
driver (Make sure the device name is correct in the `devices.yaml` and in
the corresponding simulated device file)
3. Update `tests/test_all_device_drivers.py` with the new simulated resource
11. Run the `tests/verify_physical_device_support.py` script targeting a
12. Run the `tests/verify_physical_device_support.py` script targeting a
physical device that will use the newly created driver to verify it is
working properly.

Expand Down
2 changes: 1 addition & 1 deletion src/tm_devices/helpers/enums.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -188,7 +188,7 @@ class SupportedModels(CustomStrEnum):
PSU2231A = "PSU2231A"
PSU2280 = "PSU2280"
PSU2281 = "PSU2281"
# DAQ
# DAQs
DAQ6510 = "DAQ6510"
# DMMs
DMM6500 = "DMM6500"
Expand Down
226 changes: 140 additions & 86 deletions src/tm_devices/helpers/functions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,15 +44,117 @@
####################################################################################################
# Private Constants
####################################################################################################
_KEITHLEY_2_CHAR_MODEL_LOOKUP = {
"24": "SMU",
"26": "SMU",
"22": "PSU",
"75": "DMM",
"64": "SMU",
"65": "SMU",
"37": "SS",
}
# NOTE: This regex will show as having a bunch of errors in the PyCharm IDE due to an
# open bug affecting f-strings in regex: https://youtrack.jetbrains.com/issue/PY-52140
_SUPPORTED_MODEL_REGEX_MAPPING = re.compile(
# AFGs
rf"(?P<{SupportedModels.AFG3K.value}>^AFG3\d\d\d$)"
rf"|(?P<{SupportedModels.AFG3KB.value}>^AFG3\d\d\dB$)"
rf"|(?P<{SupportedModels.AFG3KC.value}>^AFG3\d\d\dC$)"
rf"|(?P<{SupportedModels.AFG31K.value}>^AFG31\d\d\d$)"
# AWGs
rf"|(?P<{SupportedModels.AWG5200.value}>^AWG52\d\d$)"
rf"|(?P<{SupportedModels.AWG5K.value}>^AWG5\d\d\d$)"
rf"|(?P<{SupportedModels.AWG5KB.value}>^AWG5\d\d\dB$)"
rf"|(?P<{SupportedModels.AWG5KC.value}>^AWG5\d\d\dC$)"
rf"|(?P<{SupportedModels.AWG7K.value}>^AWG7\d\d\d$)"
rf"|(?P<{SupportedModels.AWG7KB.value}>^AWG7\d\d\dB$)"
rf"|(?P<{SupportedModels.AWG7KC.value}>^AWG7\d\d\dC$)"
rf"|(?P<{SupportedModels.AWG70KA.value}>^AWG70\d\d\dA$)"
rf"|(?P<{SupportedModels.AWG70KB.value}>^AWG70\d\d\dB$)"
# Scopes
rf"|(?P<{SupportedModels.DPO5K.value}>^DPO5\d\d\d$)"
rf"|(?P<{SupportedModels.DPO5KB.value}>^DPO5\d\d\dB$)"
rf"|(?P<{SupportedModels.DPO7K.value}>^DPO7\d\d\d$)"
rf"|(?P<{SupportedModels.DPO7KC.value}>^DPO7\d\d\dC$)"
rf"|(?P<{SupportedModels.DPO70K.value}>^DPO7\d\d\d\d$)"
rf"|(?P<{SupportedModels.DPO70KC.value}>^DPO7\d\d\d\dC$)"
rf"|(?P<{SupportedModels.DPO70KD.value}>^DPO7\d\d\d\dD$)"
rf"|(?P<{SupportedModels.DPO70KDX.value}>^DPO7\d\d\d\dDX$)"
rf"|(?P<{SupportedModels.DPO70KSX.value}>^DPO7\d\d\d\dSX$)"
rf"|(?P<{SupportedModels.DSA70K.value}>^DSA7\d\d\d\d$)"
rf"|(?P<{SupportedModels.DSA70KC.value}>^DSA7\d\d\d\dC$)"
rf"|(?P<{SupportedModels.DSA70KD.value}>^DSA7\d\d\d\dD$)"
rf"|(?P<{SupportedModels.LPD6.value}>^LPD6\d$)"
rf"|(?P<{SupportedModels.MSO2.value}>^MSO2\d$)"
rf"|(?P<{SupportedModels.MSO4.value}>^MSO4\d$)"
rf"|(?P<{SupportedModels.MSO4B.value}>^MSO4\dB$)"
rf"|(?P<{SupportedModels.MSO5.value}>^MSO5\d$)"
rf"|(?P<{SupportedModels.MSO5B.value}>^MSO5\dB$)"
rf"|(?P<{SupportedModels.MSO5LP.value}>^MSO5\dLP$)"
rf"|(?P<{SupportedModels.MSO6.value}>^MSO6\d$)"
rf"|(?P<{SupportedModels.MSO6B.value}>^MSO6\dB$)"
rf"|(?P<{SupportedModels.MSO2K.value}>^MSO2\d\d\d$)"
rf"|(?P<{SupportedModels.MSO2KB.value}>^MSO2\d\d\dB$)"
rf"|(?P<{SupportedModels.DPO2K.value}>^DPO2\d\d\d$)"
rf"|(?P<{SupportedModels.DPO2KB.value}>^DPO2\d\d\dB$)"
rf"|(?P<{SupportedModels.MDO3.value}>^MDO3\d$)"
rf"|(?P<{SupportedModels.MDO3K.value}>^MDO3\d\d\d$)"
rf"|(?P<{SupportedModels.MSO4K.value}>^MSO4\d\d\d$)"
rf"|(?P<{SupportedModels.MSO4KB.value}>^MSO4\d\d\dB$)"
rf"|(?P<{SupportedModels.MDO4K.value}>^MDO4\d\d\d$)"
rf"|(?P<{SupportedModels.MDO4KB.value}>^MDO4\d\d\dB$)"
rf"|(?P<{SupportedModels.MDO4KC.value}>^MDO4\d\d\dC$)"
rf"|(?P<{SupportedModels.DPO4K.value}>^DPO4\d\d\d$)"
rf"|(?P<{SupportedModels.DPO4KB.value}>^DPO4\d\d\dB$)"
rf"|(?P<{SupportedModels.MSO5K.value}>^MSO5\d\d\d$)"
rf"|(?P<{SupportedModels.MSO5KB.value}>^MSO5\d\d\dB$)"
rf"|(?P<{SupportedModels.MSO70K.value}>^MSO7\d\d\d\d$)"
rf"|(?P<{SupportedModels.MSO70KC.value}>^MSO7\d\d\d\dC$)"
rf"|(?P<{SupportedModels.MSO70KDX.value}>^MSO7\d\d\d\dDX$)"
rf"|(?P<{SupportedModels.TEKSCOPESW.value}>^TEKSCOPESW$)"
rf"|(?P<{SupportedModels.TSOVU.value}>^TSOVU$)"
rf"|(?P<{SupportedModels.TMT4.value}>^TMT4$)"
# SMUs
rf"|(?P<{SupportedModels.SMU2400.value}>^2400$)"
rf"|(?P<{SupportedModels.SMU2401.value}>^2401$)"
rf"|(?P<{SupportedModels.SMU2410.value}>^2410$)"
rf"|(?P<{SupportedModels.SMU2450.value}>^2450$)"
rf"|(?P<{SupportedModels.SMU2460.value}>^2460$)"
rf"|(?P<{SupportedModels.SMU2461.value}>^2461$)"
rf"|(?P<{SupportedModels.SMU2470.value}>^2470$)"
rf"|(?P<{SupportedModels.SMU2601B_PULSE.value}>^2601B-PULSE$)"
rf"|(?P<{SupportedModels.SMU2601B.value}>^2601B$)"
rf"|(?P<{SupportedModels.SMU2602B.value}>^2602B$)"
rf"|(?P<{SupportedModels.SMU2604B.value}>^2604B$)"
rf"|(?P<{SupportedModels.SMU2606B.value}>^2606B$)"
rf"|(?P<{SupportedModels.SMU2611B.value}>^2611B$)"
rf"|(?P<{SupportedModels.SMU2612B.value}>^2612B$)"
rf"|(?P<{SupportedModels.SMU2614B.value}>^2614B$)"
rf"|(?P<{SupportedModels.SMU2634B.value}>^2634B$)"
rf"|(?P<{SupportedModels.SMU2635B.value}>^2635B$)"
rf"|(?P<{SupportedModels.SMU2636B.value}>^2636B$)"
rf"|(?P<{SupportedModels.SMU2651A.value}>^2651A$)"
rf"|(?P<{SupportedModels.SMU2657A.value}>^2657A$)"
rf"|(?P<{SupportedModels.SMU2601A.value}>^2601A$)"
rf"|(?P<{SupportedModels.SMU2602A.value}>^2602A$)"
rf"|(?P<{SupportedModels.SMU2604A.value}>^2604A$)"
rf"|(?P<{SupportedModels.SMU2611A.value}>^2611A$)"
rf"|(?P<{SupportedModels.SMU2612A.value}>^2612A$)"
rf"|(?P<{SupportedModels.SMU2614A.value}>^2614A$)"
rf"|(?P<{SupportedModels.SMU2634A.value}>^2634A$)"
rf"|(?P<{SupportedModels.SMU2635A.value}>^2635A$)"
rf"|(?P<{SupportedModels.SMU2636A.value}>^2636A$)"
rf"|(?P<{SupportedModels.SMU6430.value}>^6430$)"
rf"|(?P<{SupportedModels.SMU6514.value}>^6514$)"
rf"|(?P<{SupportedModels.SMU6517B.value}>^6517B$)"
# PSUs
rf"|(?P<{SupportedModels.PSU2200.value}>^2200$)"
rf"|(?P<{SupportedModels.PSU2220.value}>^2220$)"
rf"|(?P<{SupportedModels.PSU2230.value}>^2230$)"
rf"|(?P<{SupportedModels.PSU2231.value}>^2231$)"
rf"|(?P<{SupportedModels.PSU2231A.value}>^2231A$)"
rf"|(?P<{SupportedModels.PSU2280.value}>^2280$)"
rf"|(?P<{SupportedModels.PSU2281.value}>^2281$)"
# DAQs
rf"|(?P<{SupportedModels.DAQ6510.value}>^DAQ6510$)"
# DMMs
rf"|(?P<{SupportedModels.DMM6500.value}>^DMM6500$)"
rf"|(?P<{SupportedModels.DMM7510.value}>^DMM7510$)"
rf"|(?P<{SupportedModels.DMM7512.value}>^DMM7512$)"
# SSs
rf"|(?P<{SupportedModels.SS3706A.value}>^3706A$)"
)


####################################################################################################
Expand Down Expand Up @@ -302,8 +404,7 @@ def detect_visa_resource_expression(input_str: str) -> Optional[Tuple[str, str]]
return retval


# pylint: disable=too-many-branches
def get_model_series(model: str) -> str: # noqa: PLR0912,C901
def get_model_series(model: str) -> str:
"""Get the series string from the full model number.
Args:
Expand All @@ -312,86 +413,39 @@ def get_model_series(model: str) -> str: # noqa: PLR0912,C901
Returns:
The model series string (ex. MSO5, LPD6).
"""
model_parts = model.strip().upper().split("-")
simplified_model = model_parts[0].replace("MODEL", "").strip()

# Remove ending characters from the model string that doesn't
# contribute to determining the correct series.
valid_model_endings = {"A", "B", "C", "D", "LP"}
model = model.strip().upper()
model_parts = model.split("-")
if len(model_parts) == 2: # noqa: PLR2004
model = model_parts[0]
elif len(model_parts) == 3: # noqa: PLR2004
model = "MODEL" + model_parts[0]

# find the model postscript if it exists and format it correctly
model_postscript = ""
model_end = ""
if re.search("[0-9]", simplified_model): # if the model contains numbers
model_end = re.split(r"\d+", simplified_model)[-1] # split on the occurrence of each number
if len(model_end) == 1 and model_end not in valid_model_endings:
simplified_model = simplified_model.rstrip(model_end)

# Check for any postscripts, e.g. 2601B-Pulse where "Pulse" is the postscript, which are a
# necessary part of determining the correct series and add them back to the model string
# to use to determine the series.
if len(model_parts) > 1:
for count, part in enumerate(model_parts):
# avoid first model part as a postscript will likely never be in the beginning
if count and len(part) > 1 and all(x.isalpha() for x in part):
model_postscript = part.capitalize()

if model.replace("MODEL", "").strip() in {x.upper() for x in SupportedModels.list_values()}:
# These are special models that don't have differences between the
# full model number and the model series.
model_beginning = SupportedModels[model.replace("MODEL", "").strip()].value
model_end = ""
else:
# Get any characters at the end of the model, e.g. 'B', 'LP', 'C'
model_beginning = ""
model_end = ""
if re.search("[0-9]", model): # if the model contains numbers
model_end = re.split(r"\d+", model)[-1] # split on the occurrence of each number
if len(model_end) == 1 and model_end not in valid_model_endings:
model_end = ""
if model_numbers := re.findall(r"\d+", model):
model_number = int(model_numbers[0])
if model.startswith("MODEL") or all(x.isdigit() for x in model.rstrip(model_end)):
model_beginning_idx = model.find(model_numbers[0])
model_beginning = model[model_beginning_idx:]
if model_beginning[-1].isalpha():
model_beginning = model_beginning[:-1]
# Convert the model into more specific driver names
with contextlib.suppress(KeyError):
model_beginning = (
_KEITHLEY_2_CHAR_MODEL_LOOKUP[model_beginning[:2]] + model_beginning
)

# This block of code can be enabled in the future if specific models need
# to be combined, e.g. SMU2601B -> SMU2600B
# Check if this is a specific, unique model series or if it needs to be a
# more generic driver.
# RELIC # if not model_postscript and (
# RELIC # any(
# RELIC # model_beginning.startswith(x)
# RELIC # for x in _MODEL_BEGINNINGS_TO_COMBINE
# RELIC # )
# RELIC # and model_beginning not in _SPECIFIC_DRIVER_SET
# RELIC # ):
# RELIC # # Models without a postscript or which are not part of a
# RELIC # # specific driver need to use a more generic driver.
# RELIC # model_beginning = model_beginning[:-2] + "00"
elif model.startswith("AWG52"):
model_beginning += model[:-1] + "0"
elif model_number >= 10_000: # noqa: PLR2004
if model.startswith("AFG"):
model_beginning += model[:5] + "K"
else:
model_beginning += model[:4] + "0K"
elif 1_000 <= model_number < 10_000: # noqa: PLR2004
model_beginning += model[:4] + "K"
else:
model_beginning += model[:4]
elif len(model_parts) >= 2: # noqa: PLR2004
model_beginning = model.capitalize()
model_beginning += model_parts[1]
else:
model_beginning = model_parts[0]

model_series = model_beginning + model_end + model_postscript

if model_series not in SupportedModels.list_values():
warnings.warn(
f'The "{model_series}" model series is not supported by {PACKAGE_NAME}', stacklevel=2
)

simplified_model += f"-{part}"

# Find the series by checking against the regex mapping
model_series = ""
if match := _SUPPORTED_MODEL_REGEX_MAPPING.match(simplified_model):
match_dict = match.groupdict()
filtered_dict = dict(filter(lambda item: item[1] is not None, match_dict.items()))
with contextlib.suppress(StopIteration):
model_series = next(iter(filtered_dict.keys()))

# Warn the user if the model is not in the regex mapping, and therefore not officially supported
if not model_series:
warnings.warn(f'The "{model}" model is not supported by {PACKAGE_NAME}', stacklevel=2)
model_series = model
return model_series


Expand Down
3 changes: 0 additions & 3 deletions tests/test_helpers.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,11 +81,9 @@ def test_create_ping_command() -> None:
("AWG5012", "AWG5K"),
("AWG70002B", "AWG70KB"),
("AFG3252", "AFG3K"),
("AFG30021", "AFG30K"),
("AFG31021", "AFG31K"),
("AFG3152C", "AFG3KC"),
("TSOVu", "TSOVu"),
("MODEL 15", "15"), # Just for complete coverage
("MODEL 2470", "SMU2470"),
("2470", "SMU2470"),
(" Model 2606B", "SMU2606B"),
Expand All @@ -101,7 +99,6 @@ def test_create_ping_command() -> None:
("DMM6500", "DMM6500"),
("MODEL 6517B", "SMU6517B"),
("2231A-30-3", "PSU2231A"),
("Tester-100", "Tester100"),
],
)
def test_get_model_series(input_string: str, expected_abbrev_model: str) -> None:
Expand Down

0 comments on commit 5369d67

Please sign in to comment.