fix overloads

dagster_polars/io_managers/delta.py

@@ -1,7 +1,7 @@
 import json
 from enum import Enum
 from pprint import pformat
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Union, overload
 
 import dagster._check as check
 import polars as pl
@@ -241,6 +241,18 @@ def write_df_to_path(
             metadata_path.parent.mkdir(parents=True, exist_ok=True)
             metadata_path.write_text(json.dumps(metadata))
 
+    @overload
+    def scan_df_from_path(
+        self, path: "UPath", context: InputContext, with_metadata: Literal[None, False]
+    ) -> pl.LazyFrame:
+        ...
+
+    @overload
+    def scan_df_from_path(
+        self, path: "UPath", context: InputContext, with_metadata: Literal[True]
+    ) -> LazyFrameWithMetadata:
+        ...
+
     def scan_df_from_path(
         self,
         path: "UPath",

dagster_polars/io_managers/parquet.py

@@ -1,5 +1,5 @@
 import json
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union, overload
 
 import polars as pl
 import pyarrow.dataset as ds
@@ -145,7 +145,7 @@ def downsteam(upstream: DataFrameWithMetadata):
                 assert metadata["my_custom_metadata"] == "my_custom_value"
     """
 
-    extension: str = ".parquet"
+    extension: str = ".parquet"  # type: ignore
 
     def sink_df_to_path(
         self,
@@ -256,12 +256,24 @@ def write_df_to_path(
                     row_group_size=row_group_size,
                 )
 
-    def scan_df_from_path(  # type: ignore
+    @overload
+    def scan_df_from_path(
+        self, path: "UPath", context: InputContext, with_metadata: Literal[None, False]
+    ) -> pl.LazyFrame:
+        ...
+
+    @overload
+    def scan_df_from_path(
+        self, path: "UPath", context: InputContext, with_metadata: Literal[True]
+    ) -> LazyFrameWithMetadata:
+        ...
+
+    def scan_df_from_path(
         self,
         path: "UPath",
         context: InputContext,
-        partition_key: Optional[str] = None,
         with_metadata: Optional[bool] = False,
+        partition_key: Optional[str] = None,
     ) -> Union[pl.LazyFrame, LazyFrameWithMetadata]:
         ldf = scan_parquet(path, context)