[CHORE] Add stubs and improve comments for pyo3-exposed abstractions,…

… + driveby type/bug fixes. (#1377)

This PR adds `.pyi` stub definitions for all of our pyo3-exposed classes
and functions, adds/improves comments for those abstractions, and makes
a bunch of driveby type/bug fixes.

The stub definitions enable stronger mypy type-checking and enable code
completion for pyo3 abstractions in IDEs.

daft/__init__.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import os
-from typing import TYPE_CHECKING
 
 from daft.logging import setup_logger
 
@@ -46,11 +45,6 @@ def get_build_type() -> str:
 
 __version__ = get_version()
 
-if TYPE_CHECKING:
-    # Placeholder for type checking for Rust module
-    class daft:
-        pass
-
 
 ###
 # Initialize analytics

daft/context.py

@@ -109,7 +109,10 @@ def logical_plan_builder_class(self) -> type[LogicalPlanBuilder]:
         from daft.logical.logical_plan import PyLogicalPlanBuilder
         from daft.logical.rust_logical_plan import RustLogicalPlanBuilder
 
-        return RustLogicalPlanBuilder if self.use_rust_planner else PyLogicalPlanBuilder
+        if self.use_rust_planner:
+            return RustLogicalPlanBuilder
+        else:
+            return PyLogicalPlanBuilder
 
 
 _DaftContext = DaftContext()