fastapi-rtk 0.2.27__py3-none-any.whl → 1.0.13__py3-none-any.whl

fastapi_rtk/types.py

@@ -1,33 +1,20 @@
-import typing
-
 import sqlalchemy.dialects.postgresql as postgresql
 from sqlalchemy import types as sa_types
 
-__all__ = ["FileColumn", "ImageColumn", "ListColumn", "JSONBListColumn"]
-
-
-class FileColumn(sa_types.TypeDecorator):
-    """
-    Extends SQLAlchemy to support and mostly identify a File Column
-    """
-
-    impl = sa_types.Text
-    cache_ok = True
-
+from .exceptions import FastAPIReactToolkitException
+from .globals import g
+from .utils import lazy
 
-class ImageColumn(sa_types.TypeDecorator):
-    """
-    Extends SQLAlchemy to support and mostly identify an Image Column
-
-    """
-
-    impl = sa_types.Text
-    cache_ok = True
-
-    def __init__(self, thumbnail_size=(20, 20, True), size=(100, 100, True), **kw):
-        sa_types.TypeDecorator.__init__(self, **kw)
-        self.thumbnail_size = thumbnail_size
-        self.size = size
+__all__ = [
+    "ListColumn",
+    "JSONBListColumn",
+    "FileColumn",
+    "ImageColumn",
+    "FileColumns",
+    "ImageColumns",
+    "JSONBFileColumns",
+    "JSONBImageColumns",
+]
 
 
 class ListColumn(sa_types.TypeDecorator):
@@ -66,4 +53,84 @@ class JSONBListColumn(ListColumn):
     impl = postgresql.JSONB
 
 
-ExportMode = typing.Literal["simplified", "detailed"]
+class BaseFileColumn:
+    manager = lazy(lambda: g.file_manager)
+    origin = "g.file_manager"
+
+    def __init__(self, allowed_extensions: list[str] | None = None, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if allowed_extensions is None:
+            allowed_extensions = self.manager.allowed_extensions
+        else:
+            # Check whether all extensions are in the allowed extensions from `g.file_manager`
+            if self.manager.allowed_extensions is not None and not all(
+                ext in self.manager.allowed_extensions for ext in allowed_extensions
+            ):
+                raise FastAPIReactToolkitException(
+                    f"Some extensions are not in the allowed extensions from {self.origin}",
+                )
+        self.allowed_extensions = allowed_extensions
+
+
+class FileColumn(BaseFileColumn, sa_types.TypeDecorator):
+    """
+    Extends SQLAlchemy to support and mostly identify a File Column.
+    """
+
+    impl = sa_types.Text
+    cache_ok = True
+
+
+class ImageColumn(BaseFileColumn, sa_types.TypeDecorator):
+    """
+    Extends SQLAlchemy to support and mostly identify an Image Column.
+    """
+
+    impl = sa_types.Text
+    cache_ok = True
+    manager = lazy(lambda: g.image_manager)
+    origin = "g.image_manager"
+
+    def __init__(
+        self,
+        allowed_extensions: list[str] | None = None,
+        thumbnail_size=(20, 20, True),
+        size=(100, 100, True),
+        *args,
+        **kwargs,
+    ):
+        super().__init__(allowed_extensions, *args, **kwargs)
+        self.thumbnail_size = thumbnail_size
+        self.size = size
+
+
+class FileColumns(FileColumn, ListColumn):
+    """
+    A column that represents a list of files.
+    """
+
+    impl = ListColumn
+
+
+class ImageColumns(ImageColumn, ListColumn):
+    """
+    A column that represents a list of images.
+    """
+
+    impl = ListColumn
+
+
+class JSONBFileColumns(FileColumn, JSONBListColumn):
+    """
+    A column that represents a list of files stored as JSONB.
+    """
+
+    impl = JSONBListColumn
+
+
+class JSONBImageColumns(ImageColumn, JSONBListColumn):
+    """
+    A column that represents a list of images stored as JSONB.
+    """
+
+    impl = JSONBListColumn

fastapi_rtk/utils/__init__.py

@@ -6,8 +6,10 @@ from .csv_json_converter import *
 from .deep_merge import *
 from .extender_mixin import *
 from .flask_appbuilder_utils import *
+from .hooks import *
 from .lazy import *
 from .merge_schema import *
+from .multiple_async_contexts import *
 from .prettify_dict import *
 from .pydantic import *
 from .run_utils import *
@@ -34,22 +36,28 @@ __all__ = [
     # .flask_appbuilder_utils
     "uuid_namegen",
     "secure_filename",
+    # .hooks
+    "hooks",
     # .lazy
     "lazy",
     "lazy_import",
     "lazy_self",
     # .merge_schema
     "merge_schema",
+    # . multiple_async_contexts
+    "multiple_async_contexts",
     # .prettify_dict
     "prettify_dict",
     # .pydantic
     "generate_schema_from_typed_dict",
+    "get_pydantic_model_field",
     # .run_utils
     "smart_run",
     "smart_run_sync",
     "safe_call",
     "safe_call_sync",
     "run_coroutine_in_threadpool",
+    "run_function_in_threadpool",
     "call_with_valid_kwargs",
     # .self_dependencies
     "SelfDepends",

fastapi_rtk/utils/async_task_runner.py

@@ -1,12 +1,16 @@
 import asyncio
 import contextvars
+import functools
 import inspect
+import traceback as tb
 import typing
 
 from .prettify_dict import prettify_dict
 
 __all__ = ["AsyncTaskRunner"]
 
+T = typing.TypeVar("T")
+
 
 class CallerInfo(typing.TypedDict):
     """
@@ -34,26 +38,37 @@ class AsyncTaskException(AsyncTaskRunnerException):
         self.caller = caller
 
 
-def wrap_in_async_task_exception(
-    task: typing.Callable[[], typing.Coroutine | None] | typing.Coroutine,
-    /,
-    caller: CallerInfo | None = None,
-):
-    async def wrapper():
-        try:
-            tsk = task
-            if callable(tsk):
-                tsk = tsk()
-            if tsk is None:
-                return None
-            return await tsk
-        except Exception as e:
-            raise AsyncTaskException(str(e), original_exception=e, caller=caller) from e
+class ResultAccessor(typing.Generic[T]):
+    """Helper class that can be awaited OR accessed directly."""
+
+    def __init__(self, task: "AsyncTask[T]"):
+        self._task = task
+
+    def __await__(self):
+        """Allow awaiting: await task.result"""
+        return self._task.__await__()
+
+    def __repr__(self) -> str:
+        """Direct access without await - returns value or raises error."""
+        return repr(self._ensure_result())
 
-    return wrapper
+    def __str__(self) -> str:
+        """Direct access without await - returns value or raises error."""
+        return str(self._ensure_result())
 
+    # Make it behave like the actual value when accessed
+    def __getattr__(self, name):
+        return getattr(self._ensure_result(), name)
 
-class AsyncTask:
+    def _ensure_result(self):
+        if not hasattr(self._task, "_result"):
+            raise RuntimeError(
+                "Task has not been executed yet. Await the task to get the result."
+            )
+        return self._task._result
+
+
+class AsyncTask(typing.Generic[T]):
     """
     Represents a task to be run asynchronously.
 
@@ -62,19 +77,51 @@ class AsyncTask:
 
     def __init__(
         self,
-        task: typing.Callable[[], typing.Coroutine | None] | typing.Coroutine,
+        task: typing.Callable[..., typing.Awaitable[T]] | typing.Awaitable[T],
         /,
         caller: CallerInfo | None = None,
         tags: list[str] | None = None,
     ):
-        self.task = wrap_in_async_task_exception(task, caller=caller)
+        self.task = task
         self.caller = caller
         self.tags = tags or []
 
+    @property
+    def result(self):
+        """
+        Provides access to the result of the task.
+
+        Returns:
+            ResultAccessor[T]: An accessor that can be awaited or accessed directly.
+        """
+        return ResultAccessor(self)
+
     def __call__(self):
-        if callable(self.task):
-            return self.task()
-        return self.task
+        return self._run()
+
+    def __await__(self):
+        return self._run().__await__()
+
+    async def _run(self):
+        """
+        Runs the task and caches the result.
+
+        Returns:
+            T: The result of the task.
+        """
+        if hasattr(self, "_result"):
+            return self._result
+
+        try:
+            coro = self.task
+            if callable(coro):
+                coro = coro()
+            self._result = await coro
+            return self._result
+        except Exception as e:
+            raise AsyncTaskException(
+                str(e), original_exception=e, caller=self.caller
+            ) from e
 
 
 class AsyncTaskRunner:
@@ -129,6 +176,9 @@ class AsyncTaskRunner:
     _task_stack: contextvars.ContextVar[typing.Optional[list[list[AsyncTask]]]] = (
         contextvars.ContextVar("_task_stack", default=None)
     )
+    _runner_stack: contextvars.ContextVar[typing.Optional[list["AsyncTaskRunner"]]] = (
+        contextvars.ContextVar("_runner_stack", default=None)
+    )
 
     def __init__(self, run_tasks_even_if_exception=False):
         """
@@ -139,23 +189,59 @@ class AsyncTaskRunner:
         """
         self.run_tasks_even_if_exception = run_tasks_even_if_exception
 
-    @staticmethod
+    @typing.overload
+    @classmethod
     def add_task(
-        *tasks: typing.Callable[[], typing.Coroutine | None] | typing.Coroutine,
+        cls,
+        task: typing.Callable[..., typing.Awaitable[T]]
+        | typing.Awaitable[T]
+        | AsyncTask[T],
+        *,
         tags: list[str] | None = None,
+        instance: "AsyncTaskRunner | None" = None,
+    ) -> AsyncTask[T]: ...
+
+    @typing.overload
+    @classmethod
+    def add_task(
+        cls,
+        task1: typing.Callable[..., typing.Awaitable[T]]
+        | typing.Awaitable[T]
+        | AsyncTask[T],
+        task2: typing.Callable[..., typing.Awaitable[T]]
+        | typing.Awaitable[T]
+        | AsyncTask[T],
+        *tasks: typing.Callable[..., typing.Awaitable[T]]
+        | typing.Awaitable[T]
+        | AsyncTask[T],
+        tags: list[str] | None = None,
+        instance: "AsyncTaskRunner | None" = None,
+    ) -> list[AsyncTask[T]]: ...
+    @classmethod
+    def add_task(
+        cls,
+        *tasks: typing.Callable[..., typing.Awaitable[T]]
+        | typing.Awaitable[T]
+        | AsyncTask[T],
+        tags: list[str] | None = None,
+        instance: "AsyncTaskRunner | None" = None,
     ):
-        stack = AsyncTaskRunner._task_stack.get()
-        if not stack:
-            raise RuntimeError(
-                "AsyncTaskRunner.add_task() called outside of context. "
-                "Use `async with AsyncTaskRunner():` before calling this."
-            )
+        """
+        Adds one or more tasks to the current context's task list. The tasks will be executed when exiting the context.
+
+        Args:
+            tags (list[str] | None, optional): Tags to associate with the tasks. Defaults to None.
+            instance (AsyncTaskRunner | None, optional): The AsyncTaskRunner instance to add tasks to. Defaults to None.
 
+        Returns:
+            AsyncTask[T] | list[AsyncTask[T]]: The added task(s).
+        """
+        task_list = cls._get_current_task_list("add_task", instance=instance)
         # Get caller info
         frame = inspect.currentframe()
         caller = inspect.getouterframes(frame, 2)[1] if frame else None
 
-        async_tasks = [
+        async_tasks: list[AsyncTask[T]] = [
             AsyncTask(
                 task,
                 caller=CallerInfo(
@@ -169,58 +255,62 @@ class AsyncTaskRunner:
             else task
             for task in tasks
         ]
-        stack[-1].extend(async_tasks)  # Add to the top context
+        task_list.extend(async_tasks)
+        return async_tasks if len(async_tasks) > 1 else async_tasks[0]
 
-    @staticmethod
-    def remove_tasks_by_tag(tag: str):
+    @classmethod
+    def remove_tasks_by_tag(
+        cls, tag: str, *, instance: "AsyncTaskRunner | None" = None
+    ):
         """
         Removes tasks with the specified tag from the current context's task list.
 
         Args:
             tag (str): The tag to filter tasks by.
         """
-        stack = AsyncTaskRunner._task_stack.get()
-        if not stack or not stack[-1]:
-            return
+        task_list = cls._get_current_task_list("remove_tasks_by_tag", instance=instance)
+        cls._update_current_task_list(
+            [task for task in task_list if tag not in task.tags], instance=instance
+        )
 
-        # Filter tasks in the current context's task list
-        stack[-1] = [task for task in stack[-1] if tag not in task.tags]
+    @classmethod
+    def get_current_runner(cls):
+        """
+        Retrieves the current AsyncTaskRunner instance from the context stack.
 
-    async def __aenter__(self):
-        stack = self._task_stack.get()
-        if stack is None:
-            stack = []
-            self._token = self._task_stack.set(stack)
-        else:
-            self._token = None  # Only reset when we set a new stack
+        Raises:
+            RuntimeError: If called outside of a context.
+
+        Returns:
+            AsyncTaskRunner: The current AsyncTaskRunner instance.
+        """
+        stack = cls._get_current_runner_stack("get_current_runner")
+        return stack[-1]
 
-        stack.append([])  # Push a new task list for this context
+    async def __aenter__(self):
+        self._handle_enter_stack()
         return self
 
     async def __aexit__(self, exc_type, exc_value, traceback):
-        stack = self._task_stack.get()
-        if not stack:
-            raise RuntimeError("Task stack corrupted or not initialized.")
-
-        tasks = stack.pop()  # Pop the current context's task list
-
-        if self._token:
-            self._task_stack.reset(self._token)
+        stack = self._get_current_stack("__aexit__")
+        tasks = self._task_list
+        stack.pop()  # Remove the current context's task list
+        self._handle_exit_stack()
 
         if exc_type and not self.run_tasks_even_if_exception:
             return
 
         if tasks:
-            exceptions: list[AsyncTaskException] = []
+            exceptions_with_index = list[tuple[AsyncTaskException, int]]()
             futures = await asyncio.gather(
                 *(task() for task in tasks), return_exceptions=True
             )
-            for future in futures:
+            for index, future in enumerate(futures):
                 if isinstance(future, AsyncTaskException):
                     # Handle exceptions from tasks
-                    exceptions.append(future)
+                    exceptions_with_index.append((future, index))
 
-            if exceptions:
+            if exceptions_with_index:
                 raise AsyncTaskRunnerException(
                     f"\n{
                         prettify_dict(
@@ -230,11 +320,146 @@ class AsyncTaskRunner:
                                     f'Task {index + 1}': {
                                         'message': str(exc),
                                         'caller': exc.caller,
-                                        'traceback': exc.original_exception.__traceback__,
+                                        'traceback': ''.join(
+                                            tb.format_exception(exc.original_exception)
+                                        ),
                                     }
-                                    for index, exc in enumerate(exceptions)
+                                    for exc, index in exceptions_with_index
                                 },
                             }
                         )
                     }"
                 )
+
+    @classmethod
+    def _get_current_task_list(
+        cls, func_name: str, *, instance: "AsyncTaskRunner | None" = None
+    ):
+        """
+        Retrieves the current task list from the context stack. If an instance is provided, it returns the task list from that instance.
+
+        Args:
+            func_name (str): The name of the function requesting the task list.
+            instance (AsyncTaskRunner | None, optional): The instance to retrieve the task list from. Defaults to None.
+
+        Raises:
+            RuntimeError: If called outside of a context.
+
+        Returns:
+            list[AsyncTask]: The current task list.
+        """
+        if instance:
+            return instance._task_list
+        return cls._get_current_stack(func_name)[-1]
+
+    @classmethod
+    def _update_current_task_list(
+        cls,
+        new_task_list: list[AsyncTask],
+        *,
+        instance: "AsyncTaskRunner | None" = None,
+    ):
+        """
+        Updates the current task list in the context stack. If an instance is provided, it updates the task list in that instance.
+
+        Args:
+            new_task_list (list[AsyncTask]): The new task list to set.
+            instance (AsyncTaskRunner | None, optional): The instance to update the task list in. Defaults to None.
+
+        Raises:
+            RuntimeError: If called outside of a context.
+        """
+        task_list = cls._get_current_task_list(
+            "update_current_task_list", instance=instance
+        )
+        task_list.clear()
+        task_list.extend(new_task_list)
+
+    @classmethod
+    def _get_current_stack(cls, func_name: str):
+        """
+        Retrieves the current task stack from the context variable.
+
+        Args:
+            func_name (str): The name of the function requesting the task stack.
+
+        Raises:
+            RuntimeError: If called outside of a context.
+
+        Returns:
+            list[list[AsyncTask]]: The current task stack.
+        """
+        stack = cls._task_stack.get()
+        if stack is None:
+            raise RuntimeError(
+                f"{cls.__name__}.{func_name}() called outside of context. "
+                f"Use `async with {cls.__name__}():` before calling this."
+            )
+        return stack
+
+    @classmethod
+    def _get_current_runner_stack(cls, func_name: str):
+        """
+        Retrieves the current runner stack from the context variable.
+
+        Args:
+            func_name (str): The name of the function requesting the runner stack.
+
+        Raises:
+            RuntimeError: If called outside of a context.
+
+        Returns:
+            list[AsyncTaskRunner]: The current runner stack.
+        """
+        stack = cls._runner_stack.get()
+        if stack is None:
+            raise RuntimeError(
+                f"{cls.__name__}.{func_name}() called outside of context. "
+                f"Use `async with {cls.__name__}():` before calling this."
+            )
+        return stack
+
+    def _handle_enter_stack(self):
+        """
+        Handles the entry into a new context by initializing or updating the task and runner stacks.
+        """
+        self._handle_init_stack()
+        stack, runner_stack = (
+            self._get_current_stack(""),
+            self._get_current_runner_stack(""),
+        )
+        self._task_list = list[AsyncTask]()
+        stack.append(self._task_list)  # Push a new task list for this context
+        runner_stack.append(self)  # Push this instance to the runner stack
+
+        # Modify instance methods to bind self
+        self.add_task = functools.partial(self.__class__.add_task, instance=self)
+        self.remove_tasks_by_tag = functools.partial(
+            self.__class__.remove_tasks_by_tag, instance=self
+        )
+
+    def _handle_exit_stack(self):
+        """
+        Cleans up the context variables for task and runner stacks if they were initialized by this instance.
+        """
+        if self._token:
+            self._task_stack.reset(self._token)
+            self._token = None
+        if self._runner_token:
+            self._runner_stack.reset(self._runner_token)
+            self._runner_token = None
+
+    def _handle_init_stack(self):
+        """
+        Ensures that the context variables for task and runner stacks are initialized.
+        """
+        try:
+            self._get_current_stack("")
+            self._token = None  # Only reset when we set a new stack
+        except RuntimeError:
+            self._token = self._task_stack.set(list[list[AsyncTask]]())
+        try:
+            self._get_current_runner_stack("")
+            self._runner_token = None  # Only reset when we set a new stack
+        except RuntimeError:
+            self._runner_token = self._runner_stack.set(list["AsyncTaskRunner"]())