From 2b2e4e0b28223ed2e1320873d19f8c5d4c57b704 Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Sat, 25 Jun 2022 15:32:20 +0200
Subject: [PATCH] Fix typing

---
 src/great_ai/great_ai/deploy/great_ai.py      | 38 ++++++++++++++++---
 .../great_ai/tracing/tracing_context.py       |  6 ++-
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/src/great_ai/great_ai/deploy/great_ai.py b/src/great_ai/great_ai/deploy/great_ai.py
index 4173047..c8fb1b8 100644
--- a/src/great_ai/great_ai/deploy/great_ai.py
+++ b/src/great_ai/great_ai/deploy/great_ai.py
@@ -9,8 +9,8 @@ from typing import (
     Optional,
     Type,
     TypeVar,
-    Union,
     cast,
+    overload,
 )
 
 from fastapi import APIRouter, FastAPI, status
@@ -45,8 +45,12 @@ class GreatAI(Generic[T]):
 
         self._func = func
 
-        def func_in_tracing_context(*args: Any, **kwargs: Any) -> Trace[T]:
-            with TracingContext[T](func.__name__) as t:
+        def func_in_tracing_context(
+            *args: Any, do_not_persist_traces: bool = False, **kwargs: Any
+        ) -> Trace[T]:
+            with TracingContext[T](
+                func.__name__, do_not_persist_traces=do_not_persist_traces
+            ) as t:
                 result = func(*args, **kwargs)
                 output = t.finalise(output=result)
             return output
@@ -68,15 +72,32 @@ class GreatAI(Generic[T]):
             redoc_url=None,
         )
 
+    @overload
     @staticmethod
-    def deploy(
+    def create(
+        func: Optional[Callable[..., T]] = None,
+    ) -> "GreatAI[T]":
+        ...
+
+    @overload
+    @staticmethod
+    def create(
+        version: str = "0.0.1",
+        disable_rest_api: bool = False,
+        disable_docs: bool = False,
+        disable_dashboard: bool = False,
+    ) -> Callable[[Callable[..., T]], "GreatAI[T]"]:
+        ...
+
+    @staticmethod
+    def create(
         func: Optional[Callable[..., T]] = None,
         *,
         version: str = "0.0.1",
         disable_rest_api: bool = False,
         disable_docs: bool = False,
         disable_dashboard: bool = False,
-    ) -> Union[Callable[[Callable[..., T]], "GreatAI[T]"], "GreatAI[T]"]:
+    ):
         if func is None:
             return cast(
                 Callable[[Callable[..., T]], GreatAI[T]],
@@ -105,9 +126,14 @@ class GreatAI(Generic[T]):
         self,
         batch: Iterable[Any],
         concurrency: Optional[int] = None,
+        do_not_persist_traces: bool = False,
     ) -> List[Trace[T]]:
         return parallel_map(
-            freeze_arguments(self._cached_func), batch, concurrency=concurrency
+            freeze_arguments(
+                partial(self._cached_func, do_not_persist_traces=do_not_persist_traces)
+            ),
+            batch,
+            concurrency=concurrency,
         )
 
     @property
diff --git a/src/great_ai/great_ai/tracing/tracing_context.py b/src/great_ai/great_ai/tracing/tracing_context.py
index a99c4d2..21bc871 100644
--- a/src/great_ai/great_ai/tracing/tracing_context.py
+++ b/src/great_ai/great_ai/tracing/tracing_context.py
@@ -24,7 +24,8 @@ T = TypeVar("T")
 class TracingContext(Generic[T]):
     _contexts: DefaultDict[int, List["TracingContext"]] = defaultdict(lambda: [])
 
-    def __init__(self, function_name: str) -> None:
+    def __init__(self, function_name: str, do_not_persist_traces: bool) -> None:
+        self._do_not_persist_traces = do_not_persist_traces
         self._models: List[Model] = []
         self._values: Dict[str, Any] = {}
         self._trace: Optional[Trace[T]] = None
@@ -90,6 +91,7 @@ class TracingContext(Generic[T]):
                 )
 
         assert self._trace is not None
-        get_context().tracing_database.save(self._trace)
+        if not self._do_not_persist_traces:
+            get_context().tracing_database.save(self._trace)
 
         return False