singlestoredb 1.14.2__cp38-abi3-win_amd64.whl → 1.15.1__cp38-abi3-win_amd64.whl

singlestoredb/docstring/util.py

@@ -0,0 +1,141 @@
+"""Utility functions for working with docstrings."""
+import typing as T
+from collections import ChainMap
+from inspect import Signature
+from itertools import chain
+
+from .common import DocstringMeta
+from .common import DocstringParam
+from .common import DocstringReturns  # noqa: F401
+from .common import DocstringStyle
+from .common import RenderingStyle
+from .parser import compose
+from .parser import parse
+
+_Func = T.Callable[..., T.Any]
+
+
+def combine_docstrings(
+    *others: _Func,
+    exclude: T.Iterable[T.Type[DocstringMeta]] = (),
+    style: DocstringStyle = DocstringStyle.AUTO,
+    rendering_style: RenderingStyle = RenderingStyle.COMPACT,
+) -> _Func:
+    """A function decorator that parses the docstrings from `others`,
+    programmatically combines them with the parsed docstring of the decorated
+    function, and replaces the docstring of the decorated function with the
+    composed result. Only parameters that are part of the decorated functions
+    signature are included in the combined docstring. When multiple sources for
+    a parameter or docstring metadata exists then the decorator will first
+    default to the wrapped function's value (when available) and otherwise use
+    the rightmost definition from ``others``.
+
+    The following example illustrates its usage:
+
+    >>> def fun1(a, b, c, d):
+    ...    '''short_description: fun1
+    ...
+    ...    :param a: fun1
+    ...    :param b: fun1
+    ...    :return: fun1
+    ...    '''
+    >>> def fun2(b, c, d, e):
+    ...    '''short_description: fun2
+    ...
+    ...    long_description: fun2
+    ...
+    ...    :param b: fun2
+    ...    :param c: fun2
+    ...    :param e: fun2
+    ...    '''
+    >>> @combine_docstrings(fun1, fun2)
+    >>> def decorated(a, b, c, d, e, f):
+    ...     '''
+    ...     :param e: decorated
+    ...     :param f: decorated
+    ...     '''
+    >>> print(decorated.__doc__)
+    short_description: fun2
+    <BLANKLINE>
+    long_description: fun2
+    <BLANKLINE>
+    :param a: fun1
+    :param b: fun1
+    :param c: fun2
+    :param e: fun2
+    :param f: decorated
+    :returns: fun1
+    >>> @combine_docstrings(fun1, fun2, exclude=[DocstringReturns])
+    >>> def decorated(a, b, c, d, e, f): pass
+    >>> print(decorated.__doc__)
+    short_description: fun2
+    <BLANKLINE>
+    long_description: fun2
+    <BLANKLINE>
+    :param a: fun1
+    :param b: fun1
+    :param c: fun2
+    :param e: fun2
+
+    :param others: callables from which to parse docstrings.
+    :param exclude: an iterable of ``DocstringMeta`` subclasses to exclude when
+        combining docstrings.
+    :param style: style composed docstring. The default will infer the style
+        from the decorated function.
+    :param rendering_style: The rendering style used to compose a docstring.
+    :return: the decorated function with a modified docstring.
+    """
+
+    def wrapper(func: _Func) -> _Func:
+        sig = Signature.from_callable(func)
+
+        comb_doc = parse(func.__doc__ or '')
+        docs = [parse(other.__doc__ or '') for other in others] + [comb_doc]
+        params = dict(
+            ChainMap(
+                *(
+                    {param.arg_name: param for param in doc.params}
+                    for doc in docs
+                ),
+            ),
+        )
+
+        for doc in reversed(docs):
+            if not doc.short_description:
+                continue
+            comb_doc.short_description = doc.short_description
+            comb_doc.blank_after_short_description = (
+                doc.blank_after_short_description
+            )
+            break
+
+        for doc in reversed(docs):
+            if not doc.long_description:
+                continue
+            comb_doc.long_description = doc.long_description
+            comb_doc.blank_after_long_description = (
+                doc.blank_after_long_description
+            )
+            break
+
+        combined: T.Dict[T.Type[DocstringMeta], T.List[DocstringMeta]] = {}
+        for doc in docs:
+            metas: T.Dict[T.Type[DocstringMeta], T.List[DocstringMeta]] = {}
+            for meta in doc.meta:
+                meta_type = type(meta)
+                if meta_type in exclude:
+                    continue
+                metas.setdefault(meta_type, []).append(meta)
+            for meta_type, meta_list in metas.items():
+                combined[meta_type] = meta_list
+
+        combined[DocstringParam] = [
+            params[name] for name in sig.parameters if name in params
+        ]
+        comb_doc.meta = list(chain(*combined.values()))
+        func.__doc__ = compose(
+            comb_doc, style=style, rendering_style=rendering_style,
+        )
+        return func
+
+    return wrapper

singlestoredb/functions/decorator.py

@@ -1,3 +1,4 @@
+import asyncio
 import functools
 import inspect
 from typing import Any
@@ -19,6 +20,7 @@ ParameterType = Union[
 ]
 
 ReturnType = ParameterType
+UDFType = Callable[..., Any]
 
 
 def is_valid_type(obj: Any) -> bool:
@@ -43,23 +45,20 @@ def is_valid_type(obj: Any) -> bool:
     return False
 
 
-def is_valid_callable(obj: Any) -> bool:
+def is_sqlstr_callable(obj: Any) -> bool:
     """Check if the object is a valid callable for a parameter type."""
     if not callable(obj):
         return False
 
     returns = utils.get_annotations(obj).get('return', None)
 
-    if inspect.isclass(returns) and issubclass(returns, str):
+    if inspect.isclass(returns) and issubclass(returns, SQLString):
         return True
 
-    raise TypeError(
-        f'callable {obj} must return a str, '
-        f'but got {returns}',
-    )
+    return False
 
 
-def expand_types(args: Any) -> Optional[Union[List[str], Type[Any]]]:
+def expand_types(args: Any) -> Optional[List[Any]]:
     """Expand the types for the function arguments / return values."""
     if args is None:
         return None
@@ -68,28 +67,32 @@ def expand_types(args: Any) -> Optional[Union[List[str], Type[Any]]]:
     if isinstance(args, str):
         return [args]
 
-    # General way of accepting pydantic.BaseModel, NamedTuple, TypedDict
-    elif is_valid_type(args):
-        return args
-
     # List of SQL strings or callables
     elif isinstance(args, list):
-        new_args = []
+        new_args: List[Any] = []
         for arg in args:
             if isinstance(arg, str):
                 new_args.append(arg)
-            elif callable(arg):
+            elif is_sqlstr_callable(arg):
                 new_args.append(arg())
+            elif type(arg) is type:
+                new_args.append(arg)
+            elif is_valid_type(arg):
+                new_args.append(arg)
             else:
                 raise TypeError(f'unrecognized type for parameter: {arg}')
         return new_args
 
     # Callable that returns a SQL string
-    elif is_valid_callable(args):
-        out = args()
-        if not isinstance(out, str):
-            raise TypeError(f'unrecognized type for parameter: {args}')
-        return [out]
+    elif is_sqlstr_callable(args):
+        return [args()]
+
+    # General way of accepting pydantic.BaseModel, NamedTuple, TypedDict
+    elif is_valid_type(args):
+        return [args]
+
+    elif type(args) is type:
+        return [args]
 
     raise TypeError(f'unrecognized type for parameter: {args}')
 
@@ -100,7 +103,8 @@ def _func(
     name: Optional[str] = None,
     args: Optional[ParameterType] = None,
     returns: Optional[ReturnType] = None,
-) -> Callable[..., Any]:
+    timeout: Optional[int] = None,
+) -> UDFType:
     """Generic wrapper for UDF and TVF decorators."""
 
     _singlestoredb_attrs = {  # type: ignore
@@ -108,6 +112,7 @@ def _func(
             name=name,
             args=expand_types(args),
             returns=expand_types(returns),
+            timeout=timeout,
         ).items() if v is not None
     }
 
@@ -115,23 +120,33 @@ def _func(
     # called later, so the wrapper much be created with the func passed
     # in at that time.
     if func is None:
-        def decorate(func: Callable[..., Any]) -> Callable[..., Any]:
-
-            def wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
-                return func(*args, **kwargs)  # type: ignore
+        def decorate(func: UDFType) -> UDFType:
 
-            wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
+            if asyncio.iscoroutinefunction(func):
+                async def async_wrapper(*args: Any, **kwargs: Any) -> UDFType:
+                    return await func(*args, **kwargs)  # type: ignore
+                async_wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
+                return functools.wraps(func)(async_wrapper)
 
-            return functools.wraps(func)(wrapper)
+            else:
+                def wrapper(*args: Any, **kwargs: Any) -> UDFType:
+                    return func(*args, **kwargs)  # type: ignore
+                wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
+                return functools.wraps(func)(wrapper)
 
         return decorate
 
-    def wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
-        return func(*args, **kwargs)  # type: ignore
-
-    wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
+    if asyncio.iscoroutinefunction(func):
+        async def async_wrapper(*args: Any, **kwargs: Any) -> UDFType:
+            return await func(*args, **kwargs)  # type: ignore
+        async_wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
+        return functools.wraps(func)(async_wrapper)
 
-    return functools.wraps(func)(wrapper)
+    else:
+        def wrapper(*args: Any, **kwargs: Any) -> UDFType:
+            return func(*args, **kwargs)  # type: ignore
+        wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
+        return functools.wraps(func)(wrapper)
 
 
 def udf(
@@ -140,7 +155,8 @@ def udf(
     name: Optional[str] = None,
     args: Optional[ParameterType] = None,
     returns: Optional[ReturnType] = None,
-) -> Callable[..., Any]:
+    timeout: Optional[int] = None,
+) -> UDFType:
     """
     Define a user-defined function (UDF).
 
@@ -167,6 +183,9 @@ def udf(
         Specifies the return data type of the function. This parameter
         works the same way as `args`. If the function is a table-valued
         function, the return type should be a `Table` object.
+    timeout : int, optional
+        The timeout in seconds for the UDF execution. If not specified,
+        the global default timeout is used.
 
     Returns
     -------
@@ -178,4 +197,5 @@ def udf(
         name=name,
         args=args,
         returns=returns,
+        timeout=timeout,
     )