ueshin commented on a change in pull request #35790:
URL: https://github.com/apache/spark/pull/35790#discussion_r826489622
##########
File path: python/pyspark/util.py
##########
@@ -421,6 +426,159 @@ def _clean_py4j_conn_for_current_thread() -> None:
thread_connection.close()
+_local = threading.local()
+
+
+def _wrap_function(class_name: str, function_name: str, func: Callable,
logger: Any) -> Callable:
+
+ signature = inspect.signature(func)
+
+ @functools.wraps(func)
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
+ if hasattr(_local, "logging") and _local.logging:
+ # no need to log since this should be internal call.
+ return func(*args, **kwargs)
+ _local.logging = True
+ try:
+ start = time.perf_counter()
+ try:
+ res = func(*args, **kwargs)
+ logger.log_success(
+ class_name, function_name, time.perf_counter() - start,
signature
+ )
+ return res
+ except Exception as ex:
+ logger.log_failure(
+ class_name, function_name, ex, time.perf_counter() -
start, signature
+ )
+ raise
+ finally:
+ _local.logging = False
+
+ return wrapper
+
+
+def _wrap_property(class_name: str, property_name: str, prop: Any, logger:
Any) -> Any:
+ @property # type: ignore[misc]
+ def wrapper(self: Any) -> Any:
+ if hasattr(_local, "logging") and _local.logging:
+ # no need to log since this should be internal call.
+ return prop.fget(self)
+ _local.logging = True
+ try:
+ start = time.perf_counter()
+ try:
+ res = prop.fget(self)
+ logger.log_success(class_name, property_name,
time.perf_counter() - start)
+ return res
+ except Exception as ex:
+ logger.log_failure(class_name, property_name, ex,
time.perf_counter() - start)
+ raise
+ finally:
+ _local.logging = False
+
+ wrapper.__doc__ = prop.__doc__
+
+ if prop.fset is not None:
+ wrapper = wrapper.setter( # type: ignore[attr-defined]
+ _wrap_function(class_name, prop.fset.__name__, prop.fset, logger)
+ )
+
+ return wrapper
+
+
+def _wrap_missing_function(
+ class_name: str, function_name: str, func: Callable, original: Any,
logger: Any
+) -> Any:
+
+ if not hasattr(original, function_name):
+ return func
+
+ signature = inspect.signature(getattr(original, function_name))
+
+ is_deprecated = func.__name__ == "deprecated_function"
+
+ @functools.wraps(func)
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
+ try:
+ return func(*args, **kwargs)
+ finally:
+ logger.log_missing(class_name, function_name, is_deprecated,
signature)
+
+ return wrapper
+
+
+def _wrap_missing_property(class_name: str, property_name: str, prop: Any,
logger: Any) -> Any:
+
+ is_deprecated = prop.fget.__name__ == "deprecated_property"
+
+ @property # type: ignore[misc]
+ def wrapper(self: Any) -> Any:
+ try:
+ return prop.fget(self)
+ finally:
+ logger.log_missing(class_name, property_name, is_deprecated)
+
+ return wrapper
+
+
+def _attach(
+ logger_module: Union[str, ModuleType], modules: Any, classes: Any,
missings: Any
Review comment:
Is it possible to avoid using `Any`s?
##########
File path: python/pyspark/util.py
##########
@@ -421,6 +426,159 @@ def _clean_py4j_conn_for_current_thread() -> None:
thread_connection.close()
+_local = threading.local()
+
+
+def _wrap_function(class_name: str, function_name: str, func: Callable,
logger: Any) -> Callable:
+
+ signature = inspect.signature(func)
+
+ @functools.wraps(func)
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
+ if hasattr(_local, "logging") and _local.logging:
+ # no need to log since this should be internal call.
+ return func(*args, **kwargs)
+ _local.logging = True
+ try:
+ start = time.perf_counter()
+ try:
+ res = func(*args, **kwargs)
+ logger.log_success(
+ class_name, function_name, time.perf_counter() - start,
signature
+ )
+ return res
+ except Exception as ex:
+ logger.log_failure(
+ class_name, function_name, ex, time.perf_counter() -
start, signature
+ )
+ raise
+ finally:
+ _local.logging = False
+
+ return wrapper
+
+
+def _wrap_property(class_name: str, property_name: str, prop: Any, logger:
Any) -> Any:
+ @property # type: ignore[misc]
+ def wrapper(self: Any) -> Any:
+ if hasattr(_local, "logging") and _local.logging:
+ # no need to log since this should be internal call.
+ return prop.fget(self)
+ _local.logging = True
+ try:
+ start = time.perf_counter()
+ try:
+ res = prop.fget(self)
+ logger.log_success(class_name, property_name,
time.perf_counter() - start)
+ return res
+ except Exception as ex:
+ logger.log_failure(class_name, property_name, ex,
time.perf_counter() - start)
+ raise
+ finally:
+ _local.logging = False
+
+ wrapper.__doc__ = prop.__doc__
+
+ if prop.fset is not None:
+ wrapper = wrapper.setter( # type: ignore[attr-defined]
+ _wrap_function(class_name, prop.fset.__name__, prop.fset, logger)
+ )
+
+ return wrapper
+
+
+def _wrap_missing_function(
+ class_name: str, function_name: str, func: Callable, original: Any,
logger: Any
+) -> Any:
+
+ if not hasattr(original, function_name):
+ return func
+
+ signature = inspect.signature(getattr(original, function_name))
+
+ is_deprecated = func.__name__ == "deprecated_function"
+
+ @functools.wraps(func)
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
+ try:
+ return func(*args, **kwargs)
+ finally:
+ logger.log_missing(class_name, function_name, is_deprecated,
signature)
+
+ return wrapper
+
+
+def _wrap_missing_property(class_name: str, property_name: str, prop: Any,
logger: Any) -> Any:
+
+ is_deprecated = prop.fget.__name__ == "deprecated_property"
+
+ @property # type: ignore[misc]
+ def wrapper(self: Any) -> Any:
+ try:
+ return prop.fget(self)
+ finally:
+ logger.log_missing(class_name, property_name, is_deprecated)
+
+ return wrapper
+
+
+def _attach(
+ logger_module: Union[str, ModuleType], modules: Any, classes: Any,
missings: Any
+) -> None:
+ if isinstance(logger_module, str):
+ logger_module = importlib.import_module(logger_module)
+
+ logger = getattr(logger_module, "get_logger")()
+
+ special_functions = set(
+ [
+ "__init__",
+ "__repr__",
+ "__str__",
+ "_repr_html_",
+ "__len__",
+ "__getitem__",
+ "__setitem__",
+ "__getattr__",
+ "__enter__",
+ "__exit__",
+ ]
+ )
+
+ # Modules
+ for target_module in modules:
+ target_name = target_module.__name__.split(".")[-1]
+ for name in getattr(target_module, "__all__"):
+ func = getattr(target_module, name)
+ if not inspect.isfunction(func):
+ continue
+ setattr(target_module, name, _wrap_function(target_name, name,
func, logger))
+
+ # Classes
+ for target_class in classes:
+ for name, func in inspect.getmembers(target_class, inspect.isfunction):
+ if name.startswith("_") and name not in special_functions:
+ continue
+ setattr(target_class, name, _wrap_function(target_class.__name__,
name, func, logger))
+
+ for name, prop in inspect.getmembers(target_class, lambda o:
isinstance(o, property)):
+ if name.startswith("_"):
+ continue
+ setattr(target_class, name, _wrap_property(target_class.__name__,
name, prop, logger))
+
+ # Missings
+ for original, missing in missings:
+ for name, func in inspect.getmembers(missing, inspect.isfunction):
+ setattr(
+ missing,
+ name,
+ _wrap_missing_function(original.__name__, name, func,
original, logger),
+ )
+
+ for name, prop in inspect.getmembers(missing, lambda o: isinstance(o,
property)):
+ setattr(missing, name, _wrap_missing_property(original.__name__,
name, prop, logger))
+
+
Review comment:
I guess we should have another file for these.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
