stubfix as cli

Author: evertlammerts

duckdb/_duckdb.pyi

@@ -694,7 +694,7 @@ class DuckDBPyRelation:
     def torch(self) -> dict: ...
     def union(self, union_rel: DuckDBPyRelation) -> DuckDBPyRelation: ...
     def unique(self, unique_aggr: str) -> DuckDBPyRelation: ...
-    def update(self, set: Expression | str, *, condition: Expression | str = None) -> None: ...
+    def update(self, set: Expression | str, *, condition: Expression | str | None = None) -> None: ...
     def value_counts(self, column: str, groups: str = "") -> DuckDBPyRelation: ...
     def var(
         self, column: str, groups: str = "", window_spec: str = "", projected_columns: str = ""

duckdb/_duckdb/functional.pyi

@@ -1,40 +1,40 @@
-import typing
+import typing as pytyping
 
 __all__: list[str] = ["ARROW", "DEFAULT", "NATIVE", "SPECIAL", "FunctionNullHandling", "PythonUDFType"]
 
 class FunctionNullHandling:
-    DEFAULT: typing.ClassVar[FunctionNullHandling]  # value = <FunctionNullHandling.DEFAULT: 0>
-    SPECIAL: typing.ClassVar[FunctionNullHandling]  # value = <FunctionNullHandling.SPECIAL: 1>
-    __members__: typing.ClassVar[
+    DEFAULT: pytyping.ClassVar[FunctionNullHandling]  # value = <FunctionNullHandling.DEFAULT: 0>
+    SPECIAL: pytyping.ClassVar[FunctionNullHandling]  # value = <FunctionNullHandling.SPECIAL: 1>
+    __members__: pytyping.ClassVar[
         dict[str, FunctionNullHandling]
     ]  # value = {'DEFAULT': <FunctionNullHandling.DEFAULT: 0>, 'SPECIAL': <FunctionNullHandling.SPECIAL: 1>}
     def __eq__(self, other: object) -> bool: ...
     def __getstate__(self) -> int: ...
     def __hash__(self) -> int: ...
     def __index__(self) -> int: ...
-    def __init__(self, value: typing.SupportsInt) -> None: ...
+    def __init__(self, value: pytyping.SupportsInt) -> None: ...
     def __int__(self) -> int: ...
     def __ne__(self, other: object) -> bool: ...
-    def __setstate__(self, state: typing.SupportsInt) -> None: ...
+    def __setstate__(self, state: pytyping.SupportsInt) -> None: ...
     @property
     def name(self) -> str: ...
     @property
     def value(self) -> int: ...
 
 class PythonUDFType:
-    ARROW: typing.ClassVar[PythonUDFType]  # value = <PythonUDFType.ARROW: 1>
-    NATIVE: typing.ClassVar[PythonUDFType]  # value = <PythonUDFType.NATIVE: 0>
-    __members__: typing.ClassVar[
+    ARROW: pytyping.ClassVar[PythonUDFType]  # value = <PythonUDFType.ARROW: 1>
+    NATIVE: pytyping.ClassVar[PythonUDFType]  # value = <PythonUDFType.NATIVE: 0>
+    __members__: pytyping.ClassVar[
         dict[str, PythonUDFType]
     ]  # value = {'NATIVE': <PythonUDFType.NATIVE: 0>, 'ARROW': <PythonUDFType.ARROW: 1>}
     def __eq__(self, other: object) -> bool: ...
     def __getstate__(self) -> int: ...
     def __hash__(self) -> int: ...
     def __index__(self) -> int: ...
-    def __init__(self, value: typing.SupportsInt) -> None: ...
+    def __init__(self, value: pytyping.SupportsInt) -> None: ...
     def __int__(self) -> int: ...
     def __ne__(self, other: object) -> bool: ...
-    def __setstate__(self, state: typing.SupportsInt) -> None: ...
+    def __setstate__(self, state: pytyping.SupportsInt) -> None: ...
     @property
     def name(self) -> str: ...
     @property

scripts/stubfix.py

@@ -232,7 +232,7 @@ def _fix_overloaded_functions(text: str) -> str:
             out.append(lines[i])
             i += 1
 
-    return "\n".join(out)
+    return "\n".join(out) + "\n"
 
 
 # -------------------------
@@ -289,15 +289,16 @@ def repl(m: re.Match) -> str:
         # Skip if annotation already explicitly optional / contains None (covers `X | None`, `Union[..., None]`, etc.)
         if "none" in ann_lower:
             return colon + ann_raw + trail
-        if ("optional[" in ann_lower) or ("typing.optional[" in ann_lower) or ("duckdb_typing.optional[" in ann_lower):
+        if "optional[" in ann_lower:
             return colon + ann_raw + trail
-        if ("classvar[" in ann_lower) or ("typing.classvar[" in ann_lower):
+        if "classvar[" in ann_lower:
             return colon + ann_raw + trail
-        if ann in {"Any", "typing.Any", "duckdb_typing.Any", "object", "ClassVar"}:
+        if ann in {"Any", "typing.Any", "pytyping.Any", "object", "ClassVar"}:
             return colon + ann_raw + trail
 
         # Otherwise wrap conservatively preserving original whitespace and annotation text
-        return f"{colon}typing.Optional[{ann}]{trail}"
+        # return f"{colon}typing.Optional[{ann}]{trail}"
+        return f"{colon}{ann} | None {trail}"
 
     return pattern.sub(repl, text)
 
@@ -348,8 +349,6 @@ def _is_valid_stubfile(path: Path) -> bool:
     )
     parser.add_argument("-r", "--recursive", action="store_true", default=False)
     args = parser.parse_args()
-    fix_stub(args.path)
-    args = parser.parse_args()
     if args.path.is_file():
         if not _is_valid_stubfile(args.path):
             print(f"{args.path} is a file, but not a pyi file.")