310 lines
9.2 KiB
Python
310 lines
9.2 KiB
Python
import inspect
|
|
import logging
|
|
import socket
|
|
import subprocess
|
|
import uuid
|
|
from contextlib import closing
|
|
from itertools import islice
|
|
from sys import version_info
|
|
|
|
PYTHON_VERSION = f"{version_info.major}.{version_info.minor}.{version_info.micro}"
|
|
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
def get_major_minor_py_version(py_version):
|
|
return ".".join(py_version.split(".")[:2])
|
|
|
|
|
|
def reraise(tp, value, tb=None):
|
|
# Taken from: https://github.com/benjaminp/six/blob/1.15.0/six.py#L694-L700
|
|
try:
|
|
if value is None:
|
|
value = tp()
|
|
if value.__traceback__ is not tb:
|
|
raise value.with_traceback(tb)
|
|
raise value
|
|
finally:
|
|
value = None
|
|
tb = None
|
|
|
|
|
|
def chunk_list(l, chunk_size):
|
|
for i in range(0, len(l), chunk_size):
|
|
yield l[i : i + chunk_size]
|
|
|
|
|
|
def _chunk_dict(d, chunk_size):
|
|
"""
|
|
Splits a dictionary into chunks of the specified size.
|
|
Taken from: https://stackoverflow.com/a/22878842
|
|
"""
|
|
it = iter(d)
|
|
for _ in range(0, len(d), chunk_size):
|
|
yield {k: d[k] for k in islice(it, chunk_size)}
|
|
|
|
|
|
def _truncate_and_ellipsize(value, max_length):
|
|
"""
|
|
Truncates the string representation of the specified value to the specified
|
|
maximum length, if necessary. The end of the string is ellipsized if truncation occurs
|
|
"""
|
|
value = str(value)
|
|
if len(value) > max_length:
|
|
return value[: (max_length - 3)] + "..."
|
|
else:
|
|
return value
|
|
|
|
|
|
def _truncate_dict(d, max_key_length=None, max_value_length=None):
|
|
"""
|
|
Truncates keys and/or values in a dictionary to the specified maximum length.
|
|
Truncated items will be converted to strings and ellipsized.
|
|
"""
|
|
key_is_none = max_key_length is None
|
|
val_is_none = max_value_length is None
|
|
|
|
if key_is_none and val_is_none:
|
|
raise ValueError("Must specify at least either `max_key_length` or `max_value_length`")
|
|
|
|
truncated = {}
|
|
for k, v in d.items():
|
|
should_truncate_key = (not key_is_none) and (len(str(k)) > max_key_length)
|
|
should_truncate_val = (not val_is_none) and (len(str(v)) > max_value_length)
|
|
|
|
new_k = _truncate_and_ellipsize(k, max_key_length) if should_truncate_key else k
|
|
if should_truncate_key:
|
|
# Use the truncated key for warning logs to avoid noisy printing to stdout
|
|
msg = f"Truncated the key `{new_k}`"
|
|
_logger.warning(msg)
|
|
|
|
new_v = _truncate_and_ellipsize(v, max_value_length) if should_truncate_val else v
|
|
if should_truncate_val:
|
|
# Use the truncated key and value for warning logs to avoid noisy printing to stdout
|
|
msg = f"Truncated the value of the key `{new_k}`. Truncated value: `{new_v}`"
|
|
_logger.warning(msg)
|
|
|
|
truncated[new_k] = new_v
|
|
|
|
return truncated
|
|
|
|
|
|
def merge_dicts(dict_a, dict_b, raise_on_duplicates=True):
|
|
"""This function takes two dictionaries and returns one singular merged dictionary.
|
|
|
|
Args:
|
|
dict_a: The first dictionary.
|
|
dict_b: The second dictionary.
|
|
raise_on_duplicates: If True, the function raises ValueError if there are duplicate keys.
|
|
Otherwise, duplicate keys in `dict_b` will override the ones in `dict_a`.
|
|
|
|
Returns:
|
|
A merged dictionary.
|
|
|
|
"""
|
|
duplicate_keys = dict_a.keys() & dict_b.keys()
|
|
if raise_on_duplicates and len(duplicate_keys) > 0:
|
|
raise ValueError(f"The two merging dictionaries contains duplicate keys: {duplicate_keys}.")
|
|
return dict_a | dict_b
|
|
|
|
|
|
def _get_fully_qualified_class_name(obj):
|
|
"""
|
|
Obtains the fully qualified class name of the given object.
|
|
"""
|
|
return obj.__class__.__module__ + "." + obj.__class__.__name__
|
|
|
|
|
|
def _inspect_original_var_name(var, fallback_name):
|
|
"""
|
|
Inspect variable name, will search above frames and fetch the same instance variable name
|
|
in the most outer frame.
|
|
If inspect failed, return fallback_name
|
|
"""
|
|
if var is None:
|
|
return fallback_name
|
|
try:
|
|
original_var_name = fallback_name
|
|
|
|
frame = inspect.currentframe().f_back
|
|
while frame is not None:
|
|
arg_info = inspect.getargvalues(frame)
|
|
|
|
fixed_args = [arg_info.locals[arg_name] for arg_name in arg_info.args]
|
|
varlen_args = list(arg_info.locals[arg_info.varargs]) if arg_info.varargs else []
|
|
keyword_args = (
|
|
list(arg_info.locals[arg_info.keywords].values()) if arg_info.keywords else []
|
|
)
|
|
|
|
all_args = fixed_args + varlen_args + keyword_args
|
|
|
|
# check whether `var` is in arg list first. If yes, go to check parent frame.
|
|
if any(var is arg for arg in all_args):
|
|
# the var is passed in from caller, check parent frame.
|
|
frame = frame.f_back
|
|
continue
|
|
|
|
for var_name, var_val in frame.f_locals.items():
|
|
if var_val is var:
|
|
original_var_name = var_name
|
|
break
|
|
|
|
break
|
|
|
|
return original_var_name
|
|
|
|
except Exception:
|
|
return fallback_name
|
|
|
|
|
|
def find_free_port():
|
|
"""
|
|
Find free socket port on local machine.
|
|
"""
|
|
with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
|
|
s.bind(("", 0))
|
|
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
return s.getsockname()[1]
|
|
|
|
|
|
def is_port_available(port: int, host: str = "127.0.0.1") -> bool:
|
|
"""
|
|
Check if a port is available for binding.
|
|
|
|
Args:
|
|
port: The port number to check.
|
|
host: The host address to check. Defaults to localhost.
|
|
|
|
Returns:
|
|
True if the port is available, False if it's already in use.
|
|
"""
|
|
with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
|
|
try:
|
|
s.bind((host, port))
|
|
return True
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def check_port_connectivity():
|
|
port = find_free_port()
|
|
try:
|
|
with subprocess.Popen(
|
|
["nc", "-l", "-p", str(port)],
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
) as server:
|
|
with subprocess.Popen(
|
|
["nc", "-zv", "localhost", str(port)],
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
) as client:
|
|
client.wait()
|
|
server.terminate()
|
|
return client.returncode == 0
|
|
except Exception as e:
|
|
_logger.warning("Failed to check port connectivity: %s", e)
|
|
return False
|
|
|
|
|
|
def is_iterator(obj):
|
|
"""
|
|
Args:
|
|
obj: Any object.
|
|
|
|
Returns:
|
|
Boolean representing whether or not 'obj' is an iterator.
|
|
"""
|
|
return (hasattr(obj, "__next__") or hasattr(obj, "next")) and hasattr(obj, "__iter__")
|
|
|
|
|
|
def _is_in_ipython_notebook():
|
|
try:
|
|
from IPython import get_ipython
|
|
|
|
return get_ipython() is not None
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def get_results_from_paginated_fn(paginated_fn, max_results_per_page, max_results=None):
|
|
"""Gets results by calling the ``paginated_fn`` until either no more results remain or
|
|
the specified ``max_results`` threshold has been reached.
|
|
|
|
Args:
|
|
paginated_fn: This function is expected to take in the number of results to retrieve
|
|
per page and a pagination token, and return a PagedList object.
|
|
max_results_per_page: The maximum number of results to retrieve per page.
|
|
max_results: The maximum number of results to retrieve overall. If unspecified,
|
|
all results will be retrieved.
|
|
|
|
Returns:
|
|
Returns a list of entities, as determined by the paginated_fn parameter, with no more
|
|
entities than specified by max_results.
|
|
|
|
"""
|
|
all_results = []
|
|
next_page_token = None
|
|
returns_all = max_results is None
|
|
while returns_all or len(all_results) < max_results:
|
|
num_to_get = max_results_per_page if returns_all else max_results - len(all_results)
|
|
if num_to_get < max_results_per_page:
|
|
page_results = paginated_fn(num_to_get, next_page_token)
|
|
else:
|
|
page_results = paginated_fn(max_results_per_page, next_page_token)
|
|
all_results.extend(page_results)
|
|
if hasattr(page_results, "token") and page_results.token:
|
|
next_page_token = page_results.token
|
|
else:
|
|
break
|
|
return all_results
|
|
|
|
|
|
class AttrDict(dict):
|
|
"""
|
|
Dict-like object that exposes its keys as attributes.
|
|
|
|
Examples
|
|
--------
|
|
>>> d = AttrDict({"a": 1, "b": 2})
|
|
>>> d.a
|
|
1
|
|
>>> d = AttrDict({"a": 1, "b": {"c": 3, "d": 4}})
|
|
>>> d.b.c
|
|
3
|
|
>>> d.c = 5
|
|
>>> d.c
|
|
5
|
|
"""
|
|
|
|
def __getattr__(self, attr):
|
|
try:
|
|
value = self[attr]
|
|
except KeyError:
|
|
raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{attr}'")
|
|
if isinstance(value, dict):
|
|
return AttrDict(value)
|
|
return value
|
|
|
|
def __setattr__(self, attr, value):
|
|
self[attr] = value
|
|
|
|
def __delattr__(self, key):
|
|
del self[key]
|
|
|
|
|
|
def get_parent_module(module):
|
|
return module[0 : module.rindex(".")]
|
|
|
|
|
|
def is_uuid(s: str) -> bool:
|
|
"""
|
|
Returns True if the specified string is a UUID, False otherwise.
|
|
"""
|
|
try:
|
|
uuid.UUID(s)
|
|
return True
|
|
except ValueError:
|
|
return False
|