refactor: simplify lazy client creation to use Type[Client] only

Address PR feedback by removing Union types and using a single consistent approach:

- Change Function/AsyncFunction constructors to accept Type[Client] only
- Remove Union[Client, Type[Client]] in favor of just Type[Client]
- Simplify _client property logic by removing isinstance checks
- Update all use() overloads to accept class types only
- Use issubclass() for async client detection instead of complex logic
- Update tests to check for _client_class attribute

This maintains the same lazy client creation behavior while being
much simpler and more consistent.

Author: zeke

src/replicate/_module_client.py

@@ -94,9 +94,11 @@ def _use(ref, *, hint=None, streaming=False, use_async=False, **kwargs):
             # For async, we need to use AsyncReplicate instead
             from ._client import AsyncReplicate
 
-            return use(lambda: AsyncReplicate(), ref, hint=hint, streaming=streaming, use_async=True, **kwargs)
+            return use(AsyncReplicate, ref, hint=hint, streaming=streaming, **kwargs)
 
-        return use(_load_client, ref, hint=hint, streaming=streaming, use_async=False, **kwargs)
+        from ._client import Replicate
+
+        return use(Replicate, ref, hint=hint, streaming=streaming, **kwargs)
 
     run = _run
     use = _use

src/replicate/lib/_predictions_use.py

@@ -9,6 +9,7 @@
     Any,
     Dict,
     List,
+    Type,
     Tuple,
     Union,
     Generic,
@@ -439,16 +440,14 @@ class Function(Generic[Input, Output]):
     _ref: str
     _streaming: bool
 
-    def __init__(self, client: Union[Client, Callable[[], Client]], ref: str, *, streaming: bool) -> None:
-        self._client_or_factory = client
+    def __init__(self, client: Type[Client], ref: str, *, streaming: bool) -> None:
+        self._client_class = client
         self._ref = ref
         self._streaming = streaming
 
     @property
     def _client(self) -> Client:
-        if callable(self._client_or_factory):
-            return self._client_or_factory()
-        return self._client_or_factory
+        return self._client_class()
 
     def __call__(self, *args: Input.args, **inputs: Input.kwargs) -> Output:
         return self.create(*args, **inputs).output()
@@ -675,16 +674,14 @@ class AsyncFunction(Generic[Input, Output]):
     _streaming: bool
     _openapi_schema: Optional[Dict[str, Any]] = None
 
-    def __init__(self, client: Union[AsyncClient, Callable[[], AsyncClient]], ref: str, *, streaming: bool) -> None:
-        self._client_or_factory = client
+    def __init__(self, client: Type[AsyncClient], ref: str, *, streaming: bool) -> None:
+        self._client_class = client
         self._ref = ref
         self._streaming = streaming
 
     @property
     def _client(self) -> AsyncClient:
-        if callable(self._client_or_factory):
-            return self._client_or_factory()
-        return self._client_or_factory
+        return self._client_class()
 
     @cached_property
     def _parsed_ref(self) -> Tuple[str, str, Optional[str]]:
@@ -814,7 +811,7 @@ async def openapi_schema(self) -> Dict[str, Any]:
 
 @overload
 def use(
-    client: Union[Client, Callable[[], Client]],
+    client: Type[Client],
     ref: Union[str, FunctionRef[Input, Output]],
     *,
     hint: Optional[Callable[Input, Output]] = None,
@@ -824,7 +821,7 @@ def use(
 
 @overload
 def use(
-    client: Union[Client, Callable[[], Client]],
+    client: Type[Client],
     ref: Union[str, FunctionRef[Input, Output]],
     *,
     hint: Optional[Callable[Input, Output]] = None,
@@ -834,7 +831,7 @@ def use(
 
 @overload
 def use(
-    client: Union[AsyncClient, Callable[[], AsyncClient]],
+    client: Type[AsyncClient],
     ref: Union[str, FunctionRef[Input, Output]],
     *,
     hint: Optional[Callable[Input, Output]] = None,
@@ -844,7 +841,7 @@ def use(
 
 @overload
 def use(
-    client: Union[AsyncClient, Callable[[], AsyncClient]],
+    client: Type[AsyncClient],
     ref: Union[str, FunctionRef[Input, Output]],
     *,
     hint: Optional[Callable[Input, Output]] = None,
@@ -853,12 +850,11 @@ def use(
 
 
 def use(
-    client: Union[Client, AsyncClient, Callable[[], Client], Callable[[], AsyncClient]],
+    client: Union[Type[Client], Type[AsyncClient]],
     ref: Union[str, FunctionRef[Input, Output]],
     *,
     hint: Optional[Callable[Input, Output]] = None,  # pylint: disable=unused-argument # noqa: ARG001 # required for type inference
     streaming: bool = False,
-    use_async: bool = False,  # Internal parameter to indicate async mode
 ) -> Union[
     Function[Input, Output],
     AsyncFunction[Input, Output],
@@ -879,12 +875,9 @@ def use(
     except AttributeError:
         pass
 
-    # Determine if this is async
-    is_async = isinstance(client, AsyncClient) or use_async
-
-    if is_async:
+    if issubclass(client, AsyncClient):
         # TODO: Fix type inference for AsyncFunction return type
-        return AsyncFunction(client, str(ref), streaming=streaming)  # type: ignore[return-value,arg-type]
-    else:
-        # TODO: Fix type inference for Function return type
-        return Function(client, str(ref), streaming=streaming)  # type: ignore[return-value,arg-type]
+        return AsyncFunction(client, str(ref), streaming=streaming)  # type: ignore[return-value]
+
+    # TODO: Fix type inference for Function return type
+    return Function(client, str(ref), streaming=streaming)  # type: ignore[return-value]

tests/test_simple_lazy.py

@@ -27,8 +27,8 @@ def test_use_does_not_create_client_immediately():
 
                 # Verify the client property is a property that will create client on demand
                 # We can't call it without a token, but we can check it's the right type
-                assert hasattr(model, "_client_or_factory")  # type: ignore[misc]
-                print("✓ Client factory is stored for lazy creation")
+                assert hasattr(model, "_client_class")  # type: ignore[misc]
+                print("✓ Client class is stored for lazy creation")
 
             except Exception as e:
                 print(f"✗ Test failed: {e}")
@@ -55,9 +55,9 @@ def test_client_created_when_model_called():
             print("✓ Model function created successfully")
 
             # Verify the model has the lazy client setup
-            assert hasattr(model, "_client_or_factory")
-            assert callable(model._client_or_factory)
-            print("✓ Lazy client factory is properly configured")
+            assert hasattr(model, "_client_class")
+            assert isinstance(model._client_class, type)
+            print("✓ Lazy client class is properly configured")
 
             # Test that accessing _client property works (creates client)
             try: