9032957364

Committed 10 May 2024 01:25PM UTC coverage: 48.538% (+1.7%) from 46.79%

Build # 9032957364

Build Type

Pull #643

github

Specific Base 74d2ec

Committed by

web-flow

Commit Message

Merge aebf2c9af into ec2d8e4fe

Pull Request Pull Request #643: 8.19.0

Run Details

377 of 1074 new or added lines in 38 files covered. (35.1%)

977 existing lines in 19 files now uncovered.

3253 of 6702 relevant lines covered (48.54%)

0.97 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

78.43

/iblrig/tools.py

import logging
import os
import re
import shutil
import socket
import subprocess
from collections.abc import Callable
from pathlib import Path
from typing import Any, TypeVar

from iblrig.constants import BONSAI_EXE
from iblrig.path_helper import create_bonsai_layout_from_template, load_pydantic_yaml
from iblrig.pydantic_definitions import RigSettings

log = logging.getLogger(__name__)


def ask_user(prompt: str, default: bool = False) -> bool:
    """
    Prompt the user for a yes/no response.

    This function displays a prompt to the user and expects a yes or no response.
    The response is not case-sensitive. If the user presses Enter without
    typing anything, the function interprets it as the default response.

    Parameters
    ----------
    prompt : str
        The prompt message to display to the user.
    default : bool, optional
        The default response when the user presses Enter without typing
        anything. If True, the default response is 'yes' (Y/y or Enter).
        If False, the default response is 'no' (N/n or Enter).

    Returns
    -------
    bool
        True if the user responds with 'yes'
        False if the user responds with 'no'
    """
    while True:
        user_input = input(f'{prompt} [Y/n] ' if default else f'{prompt} [y/N] ').strip().lower()
        if not user_input:
            return default
        elif user_input in ['y', 'yes']:
            return True
        elif user_input in ['n', 'no']:
            return False


def get_anydesk_id(silent: bool = False) -> str | None:
    """
    Retrieve the AnyDesk ID of the current machine.

    Parameters
    ----------
    silent : bool, optional
        If True, suppresses exceptions and logs them instead.
        If False (default), raises exceptions.

    Returns
    -------
    str or None
        The AnyDesk ID as a formatted string (e.g., '123 456 789') if successful,
        or None on failure.

    Raises
    ------
    FileNotFoundError
        If the AnyDesk executable is not found.
    subprocess.CalledProcessError
        If an error occurs while executing the AnyDesk command.
    StopIteration
        If the subprocess output is empty.
    UnicodeDecodeError
        If there is an issue decoding the subprocess output.

    Notes
    -----
    The function attempts to find the AnyDesk executable and retrieve the ID using the command line.
    On success, the AnyDesk ID is returned as a formatted string. If silent is True, exceptions are logged,
    and None is returned on failure. If silent is False, exceptions are raised on failure.
    """
    anydesk_id = None
    try:
        if cmd := shutil.which('anydesk'):
            pass
        elif os.name == 'nt':
            cmd = str(Path(os.environ['PROGRAMFILES(X86)'], 'AnyDesk', 'anydesk.exe'))
        if cmd is None or not Path(cmd).exists():
            raise FileNotFoundError('AnyDesk executable not found')

        proc = subprocess.Popen([cmd, '--get-id'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        if proc.stdout and re.match(r'^\d{10}$', id_string := next(proc.stdout).decode()):
            anydesk_id = f'{int(id_string):,}'.replace(',', ' ')
    except (FileNotFoundError, subprocess.CalledProcessError, StopIteration, UnicodeDecodeError) as e:
        if silent:
            log.debug(e, exc_info=True)
        else:
            raise e
    return anydesk_id


def static_vars(**kwargs) -> Callable[..., Any]:
    """
    Decorator to add static variables to a function.

    This decorator allows you to add static variables to a function by providing
    keyword arguments. Static variables are shared across all calls to the
    decorated function.

    Parameters
    ----------
    **kwargs
        Keyword arguments where the keys are variable names and the values are
        the initial values of the static variables.

    Returns
    -------
    function
        A decorated function with the specified static variables.
    """

    def decorate(func: Callable[..., Any]) -> Callable[..., Any]:
        for k in kwargs:
            setattr(func, k, kwargs[k])
        return func

    return decorate


@static_vars(return_value=None)
def internet_available(host: str = '8.8.8.8', port: int = 53, timeout: int = 3, force_update: bool = False) -> bool:
    """
    Check if the internet connection is available.

    This function checks if an internet connection is available by attempting to
    establish a connection to a specified host and port. It will use a cached
    result if the latter is available and `force_update` is set to False.

    Parameters
    ----------
    host : str, optional
        The IP address or domain name of the host to check the connection to.
        Default is "8.8.8.8" (Google's DNS server).
    port : int, optional
        The port to use for the connection check. Default is 53 (DNS port).
    timeout : int, optional
        The maximum time (in seconds) to wait for the connection attempt.
        Default is 3 seconds.
    force_update : bool, optional
        If True, force an update and recheck the internet connection even if
        the result is cached. Default is False.

    Returns
    -------
    bool
        True if an internet connection is available, False otherwise.
    """
    if not force_update and internet_available.return_value:
        return internet_available.return_value
    try:
        socket.setdefaulttimeout(timeout)
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.connect((host, port))
        internet_available.return_value = True
    except OSError:
        internet_available.return_value = False
    return internet_available.return_value


def alyx_reachable() -> bool:
    """
    Check if Alyx can be connected to.

    Returns
    -------
    bool
        True if Alyx can be connected to, False otherwise.
    """
    settings: RigSettings = load_pydantic_yaml(RigSettings)
    if settings.ALYX_URL is not None:
        return internet_available(host=settings.ALYX_URL.host, port=443, timeout=1, force_update=True)
    return False


def call_bonsai(
    workflow_file: str | Path,
    parameters: dict[str, Any] | None = None,
    start: bool = True,
    debug: bool = False,
    bootstrap: bool = True,
    editor: bool = True,
    wait: bool = True,
    check: bool = False,
) -> subprocess.Popen[bytes] | subprocess.Popen[str | bytes | Any] | subprocess.CompletedProcess:
    """
    Execute a Bonsai workflow within a subprocess call.

    Parameters
    ----------
    workflow_file : str | Path
        Path to the Bonsai workflow file.
    parameters : dict[str, str], optional
        Parameters to be passed to Bonsai workflow.
    start : bool, optional
        Start execution of the workflow within Bonsai (default is True).
    debug : bool, optional
        Enable debugging mode if True (default is False).
        Only applies if editor is True.
    bootstrap : bool, optional
        Enable Bonsai bootstrapping if True (default is True).
    editor : bool, optional
        Enable Bonsai editor if True (default is True).
    wait : bool, optional
        Wait for Bonsai process to finish (default is True).
    check : bool, optional
        Raise CalledProcessError if Bonsai process exits with non-zero exit code (default is False).
        Only applies if wait is True.

    Returns
    -------
    Popen[bytes] | Popen[str | bytes | Any] | CompletedProcess
        Pointer to the Bonsai subprocess if wait is False, otherwise subprocess.CompletedProcess.

    Raises
    ------
    FileNotFoundError
        If the Bonsai executable does not exist.
        If the specified workflow file does not exist.

    """
    if not BONSAI_EXE.exists():
        raise FileNotFoundError(BONSAI_EXE)
    workflow_file = Path(workflow_file)
    if not workflow_file.exists():
        raise FileNotFoundError(workflow_file)
    cwd = workflow_file.parent
    create_bonsai_layout_from_template(workflow_file)

    cmd = [str(BONSAI_EXE), str(workflow_file)]
    if start:
        cmd.append('--start' if debug else '--start-no-debug')
    if not editor:
        cmd.append('--no-editor')
    if not bootstrap:
        cmd.append('--no-boot')
    if parameters is not None:
        for key, value in parameters.items():
            cmd.append(f'-p:{key}={str(value)}')

    log.info(f'Starting Bonsai workflow `{workflow_file.name}`')
    log.debug(' '.join([str(x) for x in cmd]))
    if wait:
        return subprocess.run(args=cmd, cwd=cwd, check=check)
    else:
        return subprocess.Popen(args=cmd, cwd=cwd)


T = TypeVar('T', bound=object)


def get_inheritors(cls: T) -> set[T]:
    """
    Obtain a set of all direct inheritors of a class
    """
    subclasses = set(cls.__subclasses__())
    for child in subclasses:
        subclasses = subclasses.union(get_inheritors(child))
    return subclasses


class ANSI:
    """ANSI Codes for formatting text on the CLI"""

    PURPLE = '\033[95m'
    CYAN = '\033[96m'
    DARKCYAN = '\033[36m'
    BLUE = '\033[94m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'
    END = '\033[0m'

1	import logging	2✔
2	import os	2✔
3	import re	2✔
4	import shutil	2✔
5	import socket	2✔
6	import subprocess	2✔
7	from collections.abc import Callable	2✔
8	from pathlib import Path	2✔
9	from typing import Any, TypeVar	2✔
10
11	from iblrig.constants import BONSAI_EXE	2✔
12	from iblrig.path_helper import create_bonsai_layout_from_template, load_pydantic_yaml	2✔
13	from iblrig.pydantic_definitions import RigSettings	2✔
14
15	log = logging.getLogger(__name__)	2✔
16
17
18	def ask_user(prompt: str, default: bool = False) -> bool:	2✔
19	"""
20	Prompt the user for a yes/no response.
21
22	This function displays a prompt to the user and expects a yes or no response.
23	The response is not case-sensitive. If the user presses Enter without
24	typing anything, the function interprets it as the default response.
25
26	Parameters
27	----------
28	prompt : str
29	The prompt message to display to the user.
30	default : bool, optional
31	The default response when the user presses Enter without typing
32	anything. If True, the default response is 'yes' (Y/y or Enter).
33	If False, the default response is 'no' (N/n or Enter).
34
35	Returns
36	-------
37	bool
38	True if the user responds with 'yes'
39	False if the user responds with 'no'
40	"""
41	while True:	2✔
42	user_input = input(f'{prompt} [Y/n] ' if default else f'{prompt} [y/N] ').strip().lower()	2✔
43	if not user_input:	2✔
44	return default	2✔
45	elif user_input in ['y', 'yes']:	2✔
46	return True	2✔
47	elif user_input in ['n', 'no']:	2✔
48	return False	2✔
49
50
51	def get_anydesk_id(silent: bool = False) -> str \| None:	2✔
52	"""
53	Retrieve the AnyDesk ID of the current machine.
54
55	Parameters
56	----------
57	silent : bool, optional
58	If True, suppresses exceptions and logs them instead.
59	If False (default), raises exceptions.
60
61	Returns
62	-------
63	str or None
64	The AnyDesk ID as a formatted string (e.g., '123 456 789') if successful,
65	or None on failure.
66
67	Raises
68	------
69	FileNotFoundError
70	If the AnyDesk executable is not found.
71	subprocess.CalledProcessError
72	If an error occurs while executing the AnyDesk command.
73	StopIteration
74	If the subprocess output is empty.
75	UnicodeDecodeError
76	If there is an issue decoding the subprocess output.
77
78	Notes
79	-----
80	The function attempts to find the AnyDesk executable and retrieve the ID using the command line.
81	On success, the AnyDesk ID is returned as a formatted string. If silent is True, exceptions are logged,
82	and None is returned on failure. If silent is False, exceptions are raised on failure.
83	"""
UNCOV 84	anydesk_id = None	×
UNCOV 85	try:	×
UNCOV 86	if cmd := shutil.which('anydesk'):	×
UNCOV 87	pass	×
UNCOV 88	elif os.name == 'nt':	×
UNCOV 89	cmd = str(Path(os.environ['PROGRAMFILES(X86)'], 'AnyDesk', 'anydesk.exe'))	×
UNCOV 90	if cmd is None or not Path(cmd).exists():	×
UNCOV 91	raise FileNotFoundError('AnyDesk executable not found')	×
92
UNCOV 93	proc = subprocess.Popen([cmd, '--get-id'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)	×
UNCOV 94	if proc.stdout and re.match(r'^\d{10}$', id_string := next(proc.stdout).decode()):	×
UNCOV 95	anydesk_id = f'{int(id_string):,}'.replace(',', ' ')	×
UNCOV 96	except (FileNotFoundError, subprocess.CalledProcessError, StopIteration, UnicodeDecodeError) as e:	×
UNCOV 97	if silent:	×
UNCOV 98	log.debug(e, exc_info=True)	×
99	else:
UNCOV 100	raise e	×
UNCOV 101	return anydesk_id	×
102
103
104	def static_vars(**kwargs) -> Callable[..., Any]:	2✔
105	"""
106	Decorator to add static variables to a function.
107
108	This decorator allows you to add static variables to a function by providing
109	keyword arguments. Static variables are shared across all calls to the
110	decorated function.
111
112	Parameters
113	----------
114	**kwargs
115	Keyword arguments where the keys are variable names and the values are
116	the initial values of the static variables.
117
118	Returns
119	-------
120	function
121	A decorated function with the specified static variables.
122	"""
123
124	def decorate(func: Callable[..., Any]) -> Callable[..., Any]:	2✔
125	for k in kwargs:	2✔
126	setattr(func, k, kwargs[k])	2✔
127	return func	2✔
128
129	return decorate	2✔
130
131
132	@static_vars(return_value=None)	2✔
133	def internet_available(host: str = '8.8.8.8', port: int = 53, timeout: int = 3, force_update: bool = False) -> bool:	2✔
134	"""
135	Check if the internet connection is available.
136
137	This function checks if an internet connection is available by attempting to
138	establish a connection to a specified host and port. It will use a cached
139	result if the latter is available and `force_update` is set to False.
140
141	Parameters
142	----------
143	host : str, optional
144	The IP address or domain name of the host to check the connection to.
145	Default is "8.8.8.8" (Google's DNS server).
146	port : int, optional
147	The port to use for the connection check. Default is 53 (DNS port).
148	timeout : int, optional
149	The maximum time (in seconds) to wait for the connection attempt.
150	Default is 3 seconds.
151	force_update : bool, optional
152	If True, force an update and recheck the internet connection even if
153	the result is cached. Default is False.
154
155	Returns
156	-------
157	bool
158	True if an internet connection is available, False otherwise.
159	"""
160	if not force_update and internet_available.return_value:	2✔
161	return internet_available.return_value	2✔
162	try:	2✔
163	socket.setdefaulttimeout(timeout)	2✔
164	with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:	2✔
165	s.connect((host, port))	2✔
166	internet_available.return_value = True	2✔
167	except OSError:	2✔
168	internet_available.return_value = False	2✔
169	return internet_available.return_value	2✔
170
171
172	def alyx_reachable() -> bool:	2✔
173	"""
174	Check if Alyx can be connected to.
175
176	Returns
177	-------
178	bool
179	True if Alyx can be connected to, False otherwise.
180	"""
UNCOV 181	settings: RigSettings = load_pydantic_yaml(RigSettings)	×
UNCOV 182	if settings.ALYX_URL is not None:	×
UNCOV 183	return internet_available(host=settings.ALYX_URL.host, port=443, timeout=1, force_update=True)	×
UNCOV 184	return False	×
185
186
187	def call_bonsai(	2✔
188	workflow_file: str \| Path,
189	parameters: dict[str, Any] \| None = None,
190	start: bool = True,
191	debug: bool = False,
192	bootstrap: bool = True,
193	editor: bool = True,
194	wait: bool = True,
195	check: bool = False,
196	) -> subprocess.Popen[bytes] \| subprocess.Popen[str \| bytes \| Any] \| subprocess.CompletedProcess:
197	"""
198	Execute a Bonsai workflow within a subprocess call.
199
200	Parameters
201	----------
202	workflow_file : str \| Path
203	Path to the Bonsai workflow file.
204	parameters : dict[str, str], optional
205	Parameters to be passed to Bonsai workflow.
206	start : bool, optional
207	Start execution of the workflow within Bonsai (default is True).
208	debug : bool, optional
209	Enable debugging mode if True (default is False).
210	Only applies if editor is True.
211	bootstrap : bool, optional
212	Enable Bonsai bootstrapping if True (default is True).
213	editor : bool, optional
214	Enable Bonsai editor if True (default is True).
215	wait : bool, optional
216	Wait for Bonsai process to finish (default is True).
217	check : bool, optional
218	Raise CalledProcessError if Bonsai process exits with non-zero exit code (default is False).
219	Only applies if wait is True.
220
221	Returns
222	-------
223	Popen[bytes] \| Popen[str \| bytes \| Any] \| CompletedProcess
224	Pointer to the Bonsai subprocess if wait is False, otherwise subprocess.CompletedProcess.
225
226	Raises
227	------
228	FileNotFoundError
229	If the Bonsai executable does not exist.
230	If the specified workflow file does not exist.
231
232	"""
233	if not BONSAI_EXE.exists():	2✔
234	raise FileNotFoundError(BONSAI_EXE)	2✔
235	workflow_file = Path(workflow_file)	2✔
236	if not workflow_file.exists():	2✔
UNCOV 237	raise FileNotFoundError(workflow_file)	×
238	cwd = workflow_file.parent	2✔
239	create_bonsai_layout_from_template(workflow_file)	2✔
240
241	cmd = [str(BONSAI_EXE), str(workflow_file)]	2✔
242	if start:	2✔
243	cmd.append('--start' if debug else '--start-no-debug')	2✔
244	if not editor:	2✔
245	cmd.append('--no-editor')	2✔
246	if not bootstrap:	2✔
247	cmd.append('--no-boot')	2✔
248	if parameters is not None:	2✔
249	for key, value in parameters.items():	2✔
250	cmd.append(f'-p:{key}={str(value)}')	2✔
251
252	log.info(f'Starting Bonsai workflow `{workflow_file.name}`')	2✔
253	log.debug(' '.join([str(x) for x in cmd]))	2✔
254	if wait:	2✔
255	return subprocess.run(args=cmd, cwd=cwd, check=check)	2✔
256	else:
UNCOV 257	return subprocess.Popen(args=cmd, cwd=cwd)	×
258
259
260	T = TypeVar('T', bound=object)	2✔
261
262
263	def get_inheritors(cls: T) -> set[T]:	2✔
264	"""
265	Obtain a set of all direct inheritors of a class
266	"""
267	subclasses = set(cls.__subclasses__())	2✔
268	for child in subclasses:	2✔
269	subclasses = subclasses.union(get_inheritors(child))	2✔
270	return subclasses	2✔
271
272
273	class ANSI:	2✔
274	"""ANSI Codes for formatting text on the CLI"""
275
276	PURPLE = '\033[95m'	2✔
277	CYAN = '\033[96m'	2✔
278	DARKCYAN = '\033[36m'	2✔
279	BLUE = '\033[94m'	2✔
280	GREEN = '\033[92m'	2✔
281	YELLOW = '\033[93m'	2✔
282	RED = '\033[91m'	2✔
283	BOLD = '\033[1m'	2✔
284	UNDERLINE = '\033[4m'	2✔
285	END = '\033[0m'	2✔

int-brain-lab / iblrig / 9032957364

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous