update custom code

Author: saunack-shapley

src/parallel/lib/_parsing/_task_run_result.py

@@ -4,14 +4,13 @@
 
 from .._time import timeout_retry_context
 from ...types import TaskRunResult, ParsedTaskRunResult
-from ..._types import NotGiven
+from ..._types import Omit
 from ..._utils import is_str, is_given
 from ..._compat import model_parse, model_parse_json
 from ..._models import construct_type_unchecked
 from .._pydantic import is_basemodel_type
 from ._task_spec import is_output_schema_param
 from ...types.task_spec_param import OutputT, OutputSchema
-from ...types.parsed_task_run_result import ParsedTaskRunResult
 
 
 def wait_for_result(
@@ -37,7 +36,7 @@ async def wait_for_result_async(
 
 
 def task_run_result_parser(
-    run_result: TaskRunResult, output_format: Union[OutputSchema, Type[OutputT]] | None | NotGiven
+    run_result: TaskRunResult, output_format: Union[OutputSchema, Type[OutputT]] | Omit | None
 ) -> TaskRunResult | ParsedTaskRunResult[OutputT]:
     """Parse a TaskRunResult object into a ParsedTaskRunResult based on output_format."""
     if not is_given(output_format) or output_format is None or is_output_schema_param(output_format):

src/parallel/lib/_parsing/_task_spec.py

@@ -5,7 +5,7 @@
 
 import pydantic
 
-from ..._types import NOT_GIVEN, NotGiven
+from ..._types import Omit
 from ..._utils import is_str, is_dict, is_given
 from .._pydantic import to_json_schema, is_basemodel_type
 from ...types.task_spec_param import (
@@ -81,12 +81,12 @@ def _generate_output_schema(output_format: OutputSchema | Type[OutputT]) -> Outp
 
 
 def build_task_spec_param(
-    output_format: OutputSchema | Type[OutputT] | None | NotGiven,
+    output_format: OutputSchema | Type[OutputT] | Omit | None,
     _: Union[str, object],  # placeholder for input
-) -> TaskSpecParam | NotGiven:
+) -> TaskSpecParam | Omit:
     """Build a TaskSpecParam from an OutputSchema or Type[OutputT] if provided."""
     if not is_given(output_format) or output_format is None:
-        return NOT_GIVEN
+        return Omit()
 
     # output format has type OutputSchema | Type[OutputT] here
     output_schema = _generate_output_schema(output_format)  # type: ignore[arg-type]

src/parallel/resources/task_run.py

@@ -201,7 +201,7 @@ def _wait_for_result(
         *,
         run_id: str,
         deadline: float,
-        output: Optional[OutputSchema] | Type[OutputT] | NotGiven = NOT_GIVEN,
+        output: Optional[OutputSchema] | Type[OutputT] | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -229,43 +229,43 @@ def execute(
         *,
         input: Union[str, Dict[str, object]],
         processor: str,
-        metadata: Optional[Dict[str, Union[str, float, bool]]] | NotGiven = NOT_GIVEN,
-        output: Optional[OutputSchema] | NotGiven = NOT_GIVEN,
+        metadata: Optional[Dict[str, Union[str, float, bool]]] | Omit = omit,
+        output: Optional[OutputSchema] | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> TaskRunResult: ...
     @overload
     def execute(
         self,
         *,
         input: Union[str, Dict[str, object]],
         processor: str,
-        metadata: Optional[Dict[str, Union[str, float, bool]]] | NotGiven = NOT_GIVEN,
+        metadata: Optional[Dict[str, Union[str, float, bool]]] | Omit = omit,
         output: Type[OutputT],
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> ParsedTaskRunResult[OutputT]: ...
     def execute(
         self,
         *,
         input: Union[str, Dict[str, object]],
         processor: str,
-        metadata: Optional[Dict[str, Union[str, float, bool]]] | NotGiven = NOT_GIVEN,
-        output: Optional[OutputSchema] | Type[OutputT] | NotGiven = NOT_GIVEN,
+        metadata: Optional[Dict[str, Union[str, float, bool]]] | Omit = omit,
+        output: Optional[OutputSchema] | Type[OutputT] | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> TaskRunResult | ParsedTaskRunResult[OutputT]:
         """
         Convenience method to create and execute a task run in a single call.
@@ -492,7 +492,7 @@ async def _wait_for_result(
         *,
         run_id: str,
         deadline: float,
-        output: Optional[OutputSchema] | Type[OutputT] | NotGiven = NOT_GIVEN,
+        output: Optional[OutputSchema] | Type[OutputT] | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -520,39 +520,39 @@ async def execute(
         *,
         input: Union[str, Dict[str, object]],
         processor: str,
-        metadata: Optional[Dict[str, Union[str, float, bool]]] | NotGiven = NOT_GIVEN,
-        output: Optional[OutputSchema] | NotGiven = NOT_GIVEN,
+        metadata: Optional[Dict[str, Union[str, float, bool]]] | Omit = omit,
+        output: Optional[OutputSchema] | Omit = omit,
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> TaskRunResult: ...
     @overload
     async def execute(
         self,
         *,
         input: Union[str, Dict[str, object]],
         processor: str,
-        metadata: Optional[Dict[str, Union[str, float, bool]]] | NotGiven = NOT_GIVEN,
+        metadata: Optional[Dict[str, Union[str, float, bool]]] | Omit = omit,
         output: Type[OutputT],
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> ParsedTaskRunResult[OutputT]: ...
     async def execute(
         self,
         *,
         input: Union[str, Dict[str, object]],
         processor: str,
-        metadata: Optional[Dict[str, Union[str, float, bool]]] | NotGiven = NOT_GIVEN,
-        output: Optional[OutputSchema] | Type[OutputT] | NotGiven = NOT_GIVEN,
+        metadata: Optional[Dict[str, Union[str, float, bool]]] | Omit = omit,
+        output: Optional[OutputSchema] | Type[OutputT] | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> TaskRunResult | ParsedTaskRunResult[OutputT]:
         """
         Convenience method to create and execute a task run in a single call.