heyihong commented on a change in pull request #35790:
URL: https://github.com/apache/spark/pull/35790#discussion_r827059983
##########
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:
Make sense. Changed to stronger types
--
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]
