less ugly shadowing

Author: bstuder

client/lomas_client/client.py

@@ -1,7 +1,6 @@
 import base64
 import json
 import pickle
-from functools import wraps
 
 import pandas as pd
 import polars as pl
@@ -12,7 +11,6 @@
 from returns.io import IOResultE
 from returns.pipeline import flow
 from returns.pointfree import bind, map_
-from returns.unsafe import unsafe_perform_io
 
 from lomas_client.constants import (
     DUMMY_NB_ROWS,
@@ -24,7 +22,7 @@
 from lomas_client.libraries.smartnoise_sql import SmartnoiseSQLClient
 from lomas_client.libraries.smartnoise_synth import SmartnoiseSynthClient
 from lomas_client.models.config import ClientConfig
-from lomas_client.utils import parse_if_ok
+from lomas_client.utils import parse_if_ok, unwrap_all_clsmethods
 from lomas_core.constants import DPLibraries
 from lomas_core.instrumentation import init_telemetry
 from lomas_core.models.requests import GetDummyDataset, LomasRequestModel, OpenDPQueryModel
@@ -272,63 +270,20 @@ def post_processes_queries(queries: list[dict]) -> list[dict]:
         )
 
 
-# FIXME: how to cleanly shadow Client without to much python __darkmagic__ ...
-
-
-def call_and_unwrap_wrapper(method):
-    @wraps(method)
-    def call_and_unwrap(*args, **kwargs):
-        result = method(*args, **kwargs)
-        if hasattr(result, "unwrap"):
-            return unsafe_perform_io(result.unwrap())
-        return result
-
-    return call_and_unwrap
-
-
-class SmartnoiseSQLClientU(SmartnoiseSQLClient):
-    def __getattribute__(self, name, *args):
-        attr = super().__getattribute__(name)
-        if callable(attr):
-            return call_and_unwrap_wrapper(attr)
-        return attr
-
-
-class OpenDPClientU(OpenDPClient):
-    def __getattribute__(self, name, *args):
-        attr = super().__getattribute__(name)
-        if callable(attr):
-            return call_and_unwrap_wrapper(attr)
-        return attr
-
-
-class SmartnoiseSynthClientU(SmartnoiseSynthClient):
-    def __getattribute__(self, name, *args):
-        attr = super().__getattribute__(name)
-        if callable(attr):
-            return call_and_unwrap_wrapper(attr)
-        return attr
-
-
-class DiffPrivLibClientU(DiffPrivLibClient):
-    def __getattribute__(self, name, *args):
-        attr = super().__getattribute__(name)
-        if callable(attr):
-            return call_and_unwrap_wrapper(attr)
-        return attr
+@unwrap_all_clsmethods
+class Client(ClientIO):
+    """Original Client interface to shadow ClientIO whilst unwrapping results (for now)."""
 
+    def __init__(self, **kwargs: ClientConfig.model_config):
+        super().__init__(**kwargs)
 
-class Client(ClientIO):
-    def __getattribute__(self, name, *args):
-        attr = super().__getattribute__(name)
-        if callable(attr):
-            return call_and_unwrap_wrapper(attr)
-        return attr
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self.smartnoise_sql = SmartnoiseSQLClientU(self.http_client)
-        self.smartnoise_synth = SmartnoiseSynthClientU(self.http_client)
-        self.opendp = OpenDPClientU(self.http_client)
-        self.diffprivlib = DiffPrivLibClientU(self.http_client)
+        self.smartnoise_sql = unwrap_all_clsmethods(type("SmartnoiseSQLClientU", (SmartnoiseSQLClient,), {}))(
+            self.http_client
+        )
+        self.smartnoise_synth = unwrap_all_clsmethods(
+            type("SmartnoiseSynthClientU", (SmartnoiseSynthClient,), {})
+        )(self.http_client)
+        self.opendp = unwrap_all_clsmethods(type("OpenDPClientU", (OpenDPClient,), {}))(self.http_client)
+        self.diffprivlib = unwrap_all_clsmethods(type("DiffPrivLibClientU", (DiffPrivLibClient,), {}))(
+            self.http_client
+        )

client/lomas_client/utils.py

@@ -1,4 +1,8 @@
+import argparse
+import inspect
 import warnings
+from collections.abc import Callable
+from functools import wraps
 from json import JSONDecodeError
 from typing import TypeVar
 
@@ -7,6 +11,7 @@
 from pydantic import ValidationError
 from returns.functions import raise_exception
 from returns.io import IOFailure, IOResultE, IOSuccess, impure_safe
+from returns.unsafe import unsafe_perform_io
 
 from lomas_client.http_client import LomasHttpClient
 from lomas_core.constants import SSynthGanSynthesizer, SSynthMarginalSynthesizer
@@ -22,7 +27,7 @@ def parse_if_ok(res: requests.Response) -> IOResultE[str]:
     return parse_server_error(res).bind_result(specify_error_from_model)
 
 
-def parse_server_error(response: requests.Response) -> LomasServerExceptionType:
+def parse_server_error(response: requests.Response) -> IOResultE[LomasServerExceptionType]:
     """Parse a server error message based on the HTTP response.
 
     Args:
@@ -91,3 +96,27 @@ def validate_model_response(
         specify_error_from_model(job.error)
 
     return response_model.model_validate(job.result)
+
+
+def call_and_unwrap_wrapper(method: Callable) -> Callable:
+    """Unwrap IOResultE[T] back to T in the unsafest way possible."""
+
+    @wraps(method)
+    def call_and_unwrap(*args: argparse.Namespace, **kwargs: dict) -> Callable:
+        result = method(*args, **kwargs)
+        if hasattr(result, "unwrap"):
+            # First raise the internal Exception if the container is a failure
+            inner_success = result.alt(raise_exception).unwrap()
+            # Otherwise force-escape IO
+            return unsafe_perform_io(inner_success)
+        return result
+
+    return call_and_unwrap
+
+
+def unwrap_all_clsmethods(cls: type) -> type:
+    """Add a wrapper to all (public) methods of the given Class."""
+    for name, method in inspect.getmembers(cls, inspect.isfunction):
+        if not name.startswith("_"):
+            setattr(cls, name, call_and_unwrap_wrapper(method))
+    return cls