Source code for atomiq.tools.atomiq_component_creator

#! /usr/bin/env python3
"""
A helper program to create components for the components db dictionary.

.. image:: /img/tools/component_creator_1.svg
    :alt: main screen
"""
from __future__ import annotations

import atomiq
import artiq
import regex as re
import sys, inspect
import importlib
import argparse
import traceback
from pathlib import Path
import ast
from typing import Type, Callable, Optional, List, Any
import artiq.compiler.builtins as artiq_types
import artiq.experiment as aq
import readline
import pkgutil

try:
    import tomllib
except ModuleNotFoundError:
    import tomli as tomllib

from rich.console import Console
from rich.table import Table
from rich.tree import Tree
from rich.prompt import Prompt
from rich.pretty import Pretty, pretty_repr

console = Console(record=True)

ABBREVIATIONS = {
    "Photodiode": "pd",
    "LightModulator": "lm",
    "LaserSource": "laser",
    "RFSource": "rf",
    "ADCChannel": "adc",
}

COMPONENT_MODULES = ["atomiq.components"]


[docs] def import_components(modules: List) -> None: """ Imports all submodules from a module to make components inside that module visible to the component creator tool. Args: modules: List of modules, can be either strings or the modules itself """ for module in modules: if isinstance(module, str): module = importlib.import_module(module) try: for _, name, _ in pkgutil.walk_packages(module.__path__): full_name = f"{module.__name__}.{name}" try: importlib.import_module(full_name) except ModuleNotFoundError: continue import_components( [ full_name, ] ) except AttributeError: continue
[docs] def class_from_name_only(sstr: str) -> Optional[Type]: """ Inferes a type from a name string by searching in all imported modules for the class name `sstr` If multiple types are found, the most likely one is chosen based on an assesment of the parent module. If no type with the given name is found returns `None`. Args: sstr: name of the class to search for """ class_candidates = [] # Find all loaded classes with this name for module in sys.modules.values(): try: class_candidates.append(getattr(module, sstr)) except AttributeError: pass # If multiple classes are found, weigh them by different criteria likeliness = [10] * len(class_candidates) for i, class_candidate in enumerate(class_candidates): if class_candidate.__module__ == "__builtin__": # it is unlikely in atomic that a component is a buildin likeliness[i] -= 3 if atomiq.__name__ in class_candidate.__module__: likeliness[i] += 5 if artiq.__name__ in class_candidate.__module__: likeliness[i] += 4 try: if class_candidate.__module__ in sys.stdlib_module_names: likeliness[i] -= 5 except AttributeError: # stdlib_module_names is new to python3.10, just ignore for older pass return class_candidates[likeliness.index(max(likeliness))] if class_candidates else None
[docs] def class_from_full_string(sstr: str) -> Type: """ Generate a class Type from a full string. Args: sstr: a full class path like e.g. "atomiq.components.sinara.suservo.SUServoModulatedLaser" """ mod_name, class_name = sstr.rsplit(".", 1) mod = importlib.import_module(mod_name) return getattr(mod, class_name)
[docs] def get_inheritors(klass: Type) -> List[Type]: """ Recursively gets all loaded classes which are abstracting from class `klass` """ subclasses = set() work = [klass] while work: parent = work.pop() try: for child in parent.__subclasses__(): if child not in subclasses: subclasses.add(child) work.append(child) except AttributeError: # has no subclasses pass return list(subclasses)
[docs] def get_recursive_args( klass: Type, kwargs_modifications: Optional[List] = None, optionals: bool = False ) -> List: """ Analyzes the call signature of the __init__ function of `klass` to find which arguments are mandatory. This is performed recursively through all base classes. Args: kwargs_modifications: List of kwargs defined in an __init__ function which are passed down to parent classes optionals: If True return optional args, if false return mandatory args """ if kwargs_modifications is None: kwargs_modifications = [] recursive_args = [] for key, val in inspect.signature(klass).parameters.items(): if optionals: is_optional = val.default != val.empty else: is_optional = val.default == val.empty if ( is_optional and val.kind != val.VAR_KEYWORD and val.kind != val.VAR_POSITIONAL ): if key in kwargs_modifications: # kwarg item is consumed by __init__ function kwargs_modifications.pop(kwargs_modifications.index(key)) else: recursive_args.append(val) if val.kind == val.VAR_KEYWORD: # Search for **kwargs modifications kwargs_name = key.replace("*", "") kwargs_modifications.extend( [ re.findall(r"kwargs\[[\"'](.*?)[\"']\].*?=", line)[0] for line in inspect.getsourcelines(klass.__init__)[0] if re.findall(rf"{kwargs_name}\[.*?\]*=", line) ] ) init_calls = [] init_temp = "" found = False # inspect source code of __init__ functions for calls to parent __init__s for line in inspect.getsourcelines(klass.__init__)[0]: if "__init__" in line: init_temp = line.strip() found = True if init_temp.count("(") == init_temp.count(")") and found: init_calls.append(init_temp) found = False elif "__init__" not in line: init_temp += line.strip() # remove first item because it is always the __init__ itself init_calls.pop(0) # classes which should not be further inspected since they are primitives skip_classes = [ artiq.language.environment.HasEnvironment, atomiq.components.primitives.Component, ] re_initcall = re.compile(r"(?<=__init__)\((.*?)\)") for base_class in klass.__bases__: try: if base_class in skip_classes: continue # recursivels call this function on all parent classes base_class_args = get_recursive_args( base_class, kwargs_modifications, optionals ) # analyze which of the required or optional base_class_args are supplied by # kwargs or directly in __init__ calls of the child class. # The remaining args not supplied are added to recursive_args and returned. sub_init_call = [ re_initcall.search(s).group(1) for s in init_calls if base_class.__name__ in s ] if sub_init_call and not optionals: passed_args = re.split(r",(?![^[(]*[\]\)])", sub_init_call[0]) for passed_arg in passed_args: passed_arg = passed_arg.strip() if passed_arg != "self" and not (passed_arg.startswith("*")): base_class_args.pop(0) recursive_args.extend(base_class_args) except TypeError: # Is a primitive python class for example "object" pass return recursive_args
[docs] def editable_default_input(prompt: str, default: str) -> str: """ Helper function to generate a cli prompt with default input which can be edited by the user Args: prompt: Non-editable text to show default: Editable default value """ readline.set_startup_hook(lambda: readline.insert_text(default)) val = console.input(prompt) readline.set_startup_hook() return val
[docs] class ComponentEntry: """ Class representing a component configured in the components dictionary Args: name: name of the component as in the components dict target_class: type of the target class """ def __init__(self, name: str, target_class: Type): self.target_class = target_class self.class_name = target_class.__name__ self.name = name self.required_args = [ ArgumentEntry(arg, self) for arg in get_recursive_args(self.target_class) ] self.optional_args = [ ArgumentEntry(arg, self) for arg in get_recursive_args(self.target_class, optionals=True) ]
[docs] @staticmethod def init_from_template_dict( name: str, component_dict: dict[str, Any] ) -> ComponentEntry: """ Initializes a component from a *single* component dictionary entry as saved in the components database or a template. Args: name: name of the component component_dict: single component component dict entry """ config_class = ComponentEntry( name, class_from_full_string(component_dict["classname"]) ) return config_class
[docs] def populate_from_template_dict(self, component_dict: dict[str, Any]) -> None: """ Populate a component recursively from a component dictionary as saved in the components database or a template Args: component_dict: full component dict """ for key, value in component_dict[self.name]["arguments"].items(): arg_value = value if value is not None and type(value) == str: if value.startswith("&") and (value[1:] in component_dict): arg_value = ComponentEntry.init_from_template_dict( value[1:], component_dict[value[1:]] ) arg_value.populate_from_template_dict(component_dict) for defined_args in self.required_args: if defined_args.name == key: defined_args.value = arg_value break else: for defined_kwargs in self.optional_args: if defined_kwargs.name == key: defined_kwargs.value = arg_value break else: raise ValueError( f"Argument {key} defined in template is not a valid argument for {self.class_name}" )
@property def configured(self): """ Returns true if all required arguments of this and all its subcomponents are set """ return all([a.configured for a in self.required_args])
[docs] def generate_component_dict(self, component_dict: dict[str, Any]) -> None: """ Generate a component dict entry for this component and all child components recursively. The entry is directly added to the dict passed in `component_dict`. Args: component_dict: dictionary this component is added to """ arg_dict = {} for required_arg in self.get_modified_args(): if isinstance(required_arg.value, ComponentEntry): required_arg.value.generate_component_dict(component_dict) arg_dict[required_arg.name] = required_arg.get_print_value() sub_dict = { "classname": f"{self.target_class.__module__}.{self.target_class.__name__}", "arguments": arg_dict, } component_dict[self.name] = sub_dict
[docs] def get_undefined_args(self) -> List: """ Returns a list of all required arguments of this and all its subcomponents which are not yet defined/set """ undefined_args = [] for required_arg in self.required_args: if isinstance(required_arg.value, ComponentEntry): undefined_args.extend(required_arg.value.get_undefined_args()) elif required_arg.value is None: undefined_args.append(required_arg) return undefined_args
[docs] def get_fixed_args(self) -> List: """ Returns a list of all arguments of this and all its subcomponents which are defined by a fixed value (not a subcomponent) """ fixed_args = [] for fixed_arg in self.get_modified_args(): if isinstance(fixed_arg.value, ComponentEntry): fixed_args.extend(fixed_arg.value.get_fixed_args()) elif fixed_arg.value is not None: fixed_args.append(fixed_arg) return fixed_args
[docs] def get_modified_args(self) -> List: """ Returns all arguments which are either required or kwargs which differ from their default value """ return_args = [] return_args.extend(self.required_args) # force a deepcopy for opt_arg in self.optional_args: if opt_arg.value != opt_arg.default: return_args.append(opt_arg) return return_args
[docs] def set_base_name(self, old_value: str, new_value: str) -> None: """ Set the base name of this component and all its subcomponents Args: old_value: old base name used to identify the base name in the full name new_value: new base name with which the string `old_value` is replaced """ self.name = self.name.replace(old_value, new_value) for required_arg in self.required_args: if isinstance(required_arg.value, ComponentEntry): required_arg.value.set_base_name(old_value, new_value)
[docs] class ArgumentEntry: """ This class represents an argument of a component init function. Args: arg: argument of a component as given by introspection of the component parent: the component this argument is part of """ _value = None unit = None optional = False default = None def __init__(self, arg: inspect.Parameter, parent: ComponentEntry): self.raw_arg = arg self.name = arg.name if type(arg.annotation) == str: self.class_name = arg.annotation else: self.class_name = arg.annotation.__name__ self.target_class = class_from_name_only(self.class_name) self.parent = parent if arg.default != arg.empty: self.value = arg.default self.default = self.value self.optional = True @property def configured(self): """ Is true if the set value is not None or, if the value is a subcomponent, this component is fully configured """ if self.value != None: if type(self.value) == ComponentEntry: return self.value.configured else: return True else: return False @property def value(self): """ Current set value of the argument. When setting this argument, the value is automatically cast to the correct type as required by this argument" """ return self._value @value.setter def value(self, value): needs_casting = True if isinstance(value, ComponentEntry): if not issubclass(value.target_class, self.target_class): raise TypeError( "%s is not a subclass of %s required as an argument %s for %s" % ( self.target_class, value.target_class, self.name, self.parent.name, ) ) needs_casting = False elif value is None: needs_casting = False elif isinstance(value, str): if value.startswith(("&", "@")): # Link to other component/device needs_casting = False elif "*" in value: value, unit = value.split("*") value = value.strip() self.unit = unit.strip() if needs_casting: # raw values not linking to other components/devices must be castable if self.target_class == artiq_types.TFloat: value = float(value) elif self.target_class in (artiq_types.TInt32, artiq_types.TInt64): value = int(value) elif self.target_class == artiq_types.TStr: value = str(value) elif self.target_class == artiq_types.TBool: if isinstance(value, str): if value.lower() in ["true", "1", "yes"]: value = True elif value.lower() in ["false", "0", "no"]: value = False else: raise ValueError("%s is not a valid bool" % value) elif isinstance(value, int): value = bool(value) elif isinstance(value, bool): pass else: raise ValueError( "%s of type %s is not a valid bool" % (value, type(value)) ) elif self.target_class == tuple: value = tuple(value) elif self.target_class == inspect._empty: # No annotation given -> unknown type, keep as is pass else: raise NotImplementedError( "Type %s is currently not implemented" % target_class ) self._value = value
[docs] def get_print_value(self) -> PrettyValue: """ Get a nice representation of the argument value usable with the rich module and in correct representation for writing to a components dict """ return PrettyValue(self.value, self.unit)
[docs] def get_print_default(self) -> PrettyValue: """ Get a nice representation of the default argument value usable with the rich module and in correct representation for writing to a components dict """ return PrettyValue(self.default, self.unit)
[docs] class PrettyValue: """ Pretty printable value which implements a rich representation and a representation compatible with the components dict syntax """ def __init__(self, value, unit): self.value = value self.unit = unit def __repr__(self, *args): if isinstance(self.value, ComponentEntry): return f"'&{self.value.name}'" elif self.unit is not None: return f"{self.value}*{self.unit}" elif type(self.value) == str: return f"'{self.value}'" else: return str(self.value) def __rich__(self, *args): if isinstance(self.value, ComponentEntry): return f"Sub-component: [bold blue]{self.value.class_name}[/] '&{self.value.name}'" elif self.unit is not None: return f"[bold cyan]{self.value}[/]*[green]{self.unit}[/]" else: return Pretty(self.value)
[docs] def selection_gui(func: Callable) -> Callable: """ Decorator function for interactive screens allowing for going back a screen, auto clearing the console on screen change and going back a screen on error. """ def wrapper(self, *args, **kwargs): try: if not self.config["debug"]: console.clear() current = [func, args, kwargs] if self._last_screens: if self._last_screens[-1] != current: self._last_screens.append(current) else: self._last_screens.append(current) func(self, *args, **kwargs) except Exception as e: console.print(f"Something went wrong {e}") traceback.print_exc() self.go_back(index=-1) return wrapper
[docs] class InteractiveBuilder: """ Implements an interactive cli builder for configuring components and getting information on the current configuration status. The implementation uses the rich package to print to the command line Args: config: parsed atomiq tool config dictionary as defined in atomiq_tools.toml Attributes: base_name: base name string of the root component which is added to all subcomponents base_component: Type of the root component (e.q. `atomiq.components.sinara.suservo.SUServoModulatedLaser`) """ base_name: str base_component: Type def __init__(self, config: dict[str, Any]): self._last_screens = [] self.config = config
[docs] def start_from_class(self, target_class: Type, name: str) -> None: """ Start the cli tool from a single, non configured component Args: target_class: Type of the target component class name: name of the root component, used as a base name for all subcomponents """ self.base_name = name config_class = ComponentEntry(name, target_class) self.base_component = config_class self.print_class_selection(config_class)
[docs] def go_back(self, index: int = -2) -> None: """ Go back by `index` steps in the cli screen history. Args: index: determines how far to go back in screen history. The default of `-2` corresponds to the last screen, since `-1` is the current screen """ try: ls = self._last_screens[index] self._last_screens.pop(-1) selection_gui(ls[0])(self, *ls[1], **ls[2]) except IndexError: exit()
[docs] def generate_name(self, component_type: str) -> str: """ Generate a component name based on its class of the form `{class_name}_{base_name}`. If an abbreviation is found in the `config["abbreviations"]` dict, the abbreviation is used as name prefix Args: component_type: name of the component class type """ if component_type in self.config["abbreviations"]: prefix = self.config["abbreviations"][component_type] else: prefix = component_type return f"{prefix}_{self.base_name}"
[docs] def save_template(self, f_name: str) -> None: """ Save the current component dict as a template as a python file Args: f_name: Name of the output file """ component_dict = self.generate_component_dict() component_dict[self.base_component.name]["template_root_class"] = True with open(f_name, "w", encoding="utf-8") as f: f.write(pretty_repr(component_dict))
[docs] def load_template(self, f_name: str) -> None: """ Start the cli by loading and populating components from a template file Args: f_name: Name of the imput template file """ with open(f_name, "r", encoding="utf-8") as f: template_string = f.read() # needs pretreatment to make the unitfull values parseble by adding quotes around them and reading as string template_string = re.sub( r"((?:[\d.]*)\s*\*\s*(?:[\w\d.]*))", r'"\1"', template_string ) component_dict = ast.literal_eval(template_string) root_component_name, root_component_dict = next( ((k, c) for k, c in component_dict.items() if ("template_root_class" in c)), ) self.base_name = root_component_name self.base_component = ComponentEntry.init_from_template_dict( root_component_name, root_component_dict ) self.base_component.populate_from_template_dict(component_dict) self.print_class_selection(self.base_component)
[docs] def generate_component_dict(self): """ Generate a component dict by recursively gathering all configured component entries """ component_dict = {} self.base_component.generate_component_dict(component_dict) return component_dict
[docs] def get_undefined_args(self) -> list: """ Returns a list of all required arguments of all configured components which are not yet defined/set """ return self.base_component.get_undefined_args()
[docs] def get_fixed_args(self) -> list: """ Returns a list of all arguments of all configured components which are either set to a excplicit value or are optional args differing from their default value """ return self.base_component.get_fixed_args()
[docs] def set_base_name(self, value: str) -> None: """ Recursively set the base name of all configured components """ self.base_component.set_base_name(self.base_name, value) self.base_name = value
[docs] @selection_gui def print_class_selection(self, config_class: ComponentEntry) -> None: """ CLI screen for displaying and editing a component Args: config_class: Component to configure in this screen """ console.rule( f" Component [bold green]{config_class.name}[/bold green] of type [bold cyan]{config_class.class_name}[/bold cyan] requires the following sub components:", align="left", ) if link := generate_atomiq_class_doc_link(config_class.target_class): console.print( f"More information about this component can be found in the {link}." ) all_args_sorted_table = [] # all arguments in the order given in the table # Generate table with mandatory arguments arg_table = Table() arg_table.add_column("№") arg_table.add_column("Argument") arg_table.add_column("Class") arg_table.add_column("Value", max_width=55) arg_table.add_column("Complete") for i, required_arg in enumerate(config_class.required_args): all_args_sorted_table.append(required_arg) arg_table.add_row( Pretty(i), required_arg.name, required_arg.class_name, required_arg.get_print_value(), Pretty(required_arg.configured), ) console.print(arg_table) # Generate table with optional arguments which already differ from their default n_required_args = len(config_class.required_args) defined_opts_table = Table() defined_opts_table.add_column("№") defined_opts_table.add_column("Argument") defined_opts_table.add_column("Value") defined_opts_table.add_column("Default") defined_opts_table.add_column("Type") i_defined_opts = 0 for optional_arg in config_class.optional_args: if optional_arg.value != optional_arg.default: all_args_sorted_table.append(optional_arg) defined_opts_table.add_row( Pretty(i_defined_opts + n_required_args), optional_arg.name, optional_arg.get_print_value(), optional_arg.get_print_default(), optional_arg.class_name, ) i_defined_opts += 1 if i_defined_opts > 0: console.rule( f" The following [bold green]optional[/bold green] arguments are defined:", align="left", ) console.print(defined_opts_table) console.rule() # Generate table with additional optional arguments console.rule( f" The following [bold green]optional[/bold green] arguments can be defined:", align="left", ) opts_table = Table() opts_table.add_column("№") opts_table.add_column("Argument") opts_table.add_column("Value") opts_table.add_column("Type") for optional_arg in config_class.optional_args: if optional_arg.value == optional_arg.default: all_args_sorted_table.append(optional_arg) opts_table.add_row( Pretty(i_defined_opts + n_required_args), optional_arg.name, optional_arg.get_print_value(), optional_arg.class_name, ) i_defined_opts += 1 console.print(opts_table) console.rule() command_table = Table("", "", box=None) command_table.add_row("\[num]", ": Edit entry") command_table.add_row("v\[num]", ": Change component value") command_table.add_row("c\[num]", ": Add/change subcomponent from list") command_table.add_row("u", ": To show undefined arguments") command_table.add_row("f", ": To show fixed value arguments") command_table.add_row("n", ": Change root component name") command_table.add_row("t", ": Show dependency tree") command_table.add_row("c", ": To show current component dict") command_table.add_row("st", ": To save template") command_table.add_row("b", ": Go back") console.print(command_table) command = console.input(": ") if command == "b": self.go_back() elif command == "t": self.tree_view() elif command == "u": self.print_undefined_args() elif command == "f": self.print_fixed_args() elif command == "st": template_name = editable_default_input( "Enter template name: ", "template.py" ) readline.set_startup_hook() self.save_template(template_name) command = console.input("Template saved. Press any key to go back...") self.go_back(-1) elif command == "c": console.print(pretty_repr(self.generate_component_dict())) command = console.input("Press any key to go back...") self.go_back(-1) elif command == "n": command = Prompt.ask( "Enter a new root component name: ", default=self.base_name ) self.set_base_name(command) self.go_back(-1) elif command == "print_screen": console.save_svg("tool_output.svg") elif command.startswith("v"): new_name = console.input("Enter component value:") arg_selected = all_args_sorted_table[int(command[1:])] arg_selected.value = new_name self.go_back(-1) elif command.startswith("c"): arg_selected = all_args_sorted_table[int(command[1:])] self.print_arg_possibilities(arg_selected) else: arg_selected = all_args_sorted_table[int(command)] if arg_selected.value is None: self.print_arg_possibilities(arg_selected) elif isinstance(arg_selected.value, ComponentEntry): self.print_class_selection(arg_selected.value) else: new_name = console.input("Enter component value:") arg_selected.value = new_name self.go_back(-1)
[docs] @selection_gui def print_arg_possibilities(self, arg: ArgumentEntry) -> None: """ CLI screen for displaying and choosing from possible components for a given argument of another component. Possible components are all classes which inherit from the type of the argument entry. All loaded modules are searched for possible matches. Args: ArgumentEntry: Argument for which a component should be chosen. """ console.rule( f"The following options for [bold green]{arg.name}[/bold green] of type [bold cyan]{arg.class_name}[/bold cyan] are available:", align="left", ) arg_table = Table() arg_table.add_column("№") arg_table.add_column("Class") arg_table.add_column("Documentation") target_class = class_from_name_only(arg.class_name) inheritors = get_inheritors(target_class) inheritors.append( target_class ) # We want the base class itself also as an option inheritors.sort(key=lambda x: x.__name__) for i, inheritor in enumerate(inheritors): arg_table.add_row( str(i), inheritor.__name__, generate_atomiq_class_doc_link(inheritor), ) console.print(arg_table) command = console.input("Select a target component to continue or go \[b]ack: ") if command == "b": self.go_back() else: arg.value = ComponentEntry( self.generate_name(arg.class_name), inheritors[int(command)], ) self.print_class_selection(arg.value)
[docs] @selection_gui def tree_view(self) -> None: """ CLI screen which shows the relations of the currently configured components in a tree-style diagram """ def build_tree(component, tree): sub_tree = tree.add(f"{component.name} ({component.class_name})") for sub_comp in component.required_args: if sub_comp.value is not None: if isinstance(sub_comp.value, ComponentEntry): build_tree(sub_comp.value, sub_tree) elif isinstance(sub_comp.value, str): if sub_comp.value.startswith(("&", "%")): sub_tree.add(sub_comp.value) tree = Tree("Component Tree") build_tree(self.base_component, tree) console.print(tree) console.input("Press any key to go back") self.go_back()
[docs] @selection_gui def print_undefined_args(self) -> None: """ CLI screen to display all required args which are not (yet) set/defined """ console.rule( f"Undefined Arguments:", align="left", ) arg_table = Table() arg_table.add_column("№") arg_table.add_column("Parent") arg_table.add_column("Argument") arg_table.add_column("Class") undefined_args = self.get_undefined_args() for i, required_arg in enumerate(undefined_args): arg_table.add_row( pretty_repr(i), required_arg.parent.name, required_arg.name, required_arg.class_name, ) console.print(arg_table) console.rule() command_table = Table("", "", box=None) command_table.add_row("\[num]", ": Add component") command_table.add_row("v\[num]", ": Change component value") command_table.add_row("b", ": Go back") console.print(command_table) command = console.input(": ") if command == "b": self.go_back() elif command.startswith("v"): new_name = console.input("Enter component value:") undefined_args[int(command[1])].value = new_name self.go_back(-1) else: self.print_arg_possibilities(undefined_args[int(command)])
[docs] @selection_gui def print_fixed_args(self): """ CLI screen to display all args which are set to a excplicit value (not another component) """ console.rule( f"Arguments with fixed values:", align="left", ) arg_table = Table() arg_table.add_column("№") arg_table.add_column("Parent") arg_table.add_column("Argument") arg_table.add_column("Value") fixed_args = self.get_fixed_args() for i, required_arg in enumerate(fixed_args): arg_table.add_row( Pretty(i), required_arg.parent.name, required_arg.name, required_arg.get_print_value(), ) console.print(arg_table) console.rule() command_table = Table("", "", box=None) command_table.add_row("\[num]", ": Change component value") command_table.add_row("b", ": Go back") console.print(command_table) command = console.input(": ") if command == "b": self.go_back() else: new_name = editable_default_input( "Enter component value: ", fixed_args[int(command)].value ) readline.set_startup_hook() fixed_args[int(command)].value = new_name self.go_back(-1)
[docs] def get_argparser() -> ArgumentParser: common_parser = argparse.ArgumentParser(add_help=False) common_parser.add_argument( "-A", "--abbreviations", default=None, help=f"Abbreviation dictionary. Default is {ABBREVIATIONS}. Overrides dict given in a config file (-c)", ) common_parser.add_argument( "-c", "--config-file", default=None, help="Specify a config file in toml format. Searches for 'atomiq_tools.toml' by default", ) common_parser.add_argument( "-m", "--component-modules", default=None, nargs="+", help="Specify additional modules to search for components. Default is [atomiq.components]", ) common_parser.add_argument( "-d", "--debug", action="store_true", help="Enable debug mode", ) parser = argparse.ArgumentParser( prog=Path(__file__).stem, description="A helper program to create component db device chains", parents=[common_parser], ) subparsers = parser.add_subparsers(dest="command") p_template = subparsers.add_parser( "template", help="create component chain from template", parents=[common_parser] ) p_template.add_argument("filename") p_class = subparsers.add_parser( "class", help="start new component chain from a target component class", parents=[common_parser], ) p_class.add_argument( "comp_class", help="Full path of the target component class (e.g. atomiq.components.sinara.suservo.SUServoModulatedLaser)", ) p_class.add_argument("name", help="Name of the component") return parser
[docs] def main(): parser = get_argparser() args = parser.parse_args() config_dict = { "abbreviations": ABBREVIATIONS, "component_modules": COMPONENT_MODULES, } if Path("atomiq_tools.toml").is_file(): with open("atomiq_tools.toml", "rb") as f: config_dict = config_dict | tomllib.load(f) if args.config_file is not None: with open(args.config_file, "rb") as f: config_dict = config_dict | tomllib.load(f) if args.abbreviations is not None: config_dict["abbreviations"] = args.abbreviations if args.component_modules is not None: config_dict["component_modules"].extend(args.component_modules) config_dict["debug"] = args.debug import_components(config_dict["component_modules"]) builder = InteractiveBuilder(config_dict) if args.command == "template": builder.load_template(args.filename) elif args.command == "class": builder.start_from_class(class_from_full_string(args.comp_class), args.name) else: parser.print_help()
if __name__ == "__main__": try: main() except KeyboardInterrupt: print("Exiting...")