Files
ophyd_devices/.github/scripts/retrieve_device_classes.py

98 lines
3.7 KiB
Python

"""
This module contains helper methods to retrieve a list of all devices in a directory.
It does so by recursively traversing the directory and checking if the file contains a
class that inherits from the Device class.
"""
import importlib.util
import os
from ophyd import Device
def get_devices(repo_name: str, ignore: str = "") -> dict:
"""
Get all devices in a directory.
Args:
directory (str): The directory to search for devices.
ignore (str): A comma-separated list of module names to ignore.
Returns:
list: A list of all devices in the directory.
"""
devices = {}
anchor_module = importlib.import_module(f"{repo_name}.devices")
directory = os.path.dirname(anchor_module.__file__)
for root, _, files in os.walk(directory):
for file in files:
if not file.endswith(".py") or file.startswith("__"):
continue
path = os.path.join(root, file)
subs = os.path.dirname(os.path.relpath(path, directory)).split("/")
if len(subs) == 1 and not subs[0]:
module_name = file.split(".")[0]
else:
module_name = ".".join(subs + [file.split(".")[0]])
if module_name in ignore.split(","):
continue
module = importlib.import_module(f"{repo_name}.devices.{module_name}")
for name in dir(module):
obj = getattr(module, name)
if not hasattr(obj, "__module__") or obj.__module__ != module.__name__:
continue
if isinstance(obj, type) and issubclass(obj, Device) and obj is not Device:
devices[obj.__name__] = obj
return dict(sorted(devices.items(), key=lambda x: x[0].lower()))
def write_device_list(devices: dict, file_name: str, repo_name: str, append=False):
"""
Write the list of devices to a file.
Args:
devices (list): A list of devices.
file_out (str): The file to write the devices to.
repo_name (str): The repository name for linking to the source code.
append (bool): Whether to append to the file or overwrite it.
"""
if not append or not os.path.exists(file_name):
with open(file_name, "w", encoding="utf-8") as output_file:
output_file.write("// This file was autogenerated. Do not edit it manually.\n")
output_file.write("## Device List\n")
with open(file_name, "a", encoding="utf-8") as output_file:
output_file.write(f"### {repo_name} \n")
output_file.write("| Device | Documentation | Module |\n")
output_file.write("| :----- | :------------- | :------ |\n")
for dev, cls in devices.items():
doc = cls.__doc__
doc = "" if doc is None else doc.replace("\n", "<br>")
output_file.write(
f"| {dev} | {doc} | [{cls.__module__}](https://github.com/bec-project/{repo_name}/tree/main/{cls.__module__.replace('.', '/')}.py) |\n"
)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Retrieve a list of devices in a directory.")
parser.add_argument("repo", type=str, help="The repository to link to.")
parser.add_argument("output", type=str, help="The file to write the devices to.")
parser.add_argument(
"--append", action="store_true", help="Append to the file instead of overwriting it."
)
parser.add_argument(
"--ignore", type=str, help="Ignore the specified modules (comma-separated).", default=""
)
args = parser.parse_args()
devs = get_devices(args.repo, ignore=args.ignore)
write_device_list(devs, args.output, args.repo, append=args.append)