From ae9da649e2ad05be22b57ee4761d83e02fb88376 Mon Sep 17 00:00:00 2001
From: chjuncn <130543538+chjuncn@users.noreply.github.com>
Date: Thu, 23 Jan 2025 12:07:38 +0800
Subject: [PATCH] Logical.py and its ancestor accept certain params and lint
 refactor

lint refactor.

Logical.py and its ancestor accept very certain params, the change will make this clear.
---
 src/palimpzest/query/generators/generators.py          |  4 ++--
 .../query/operators/code_synthesis_convert.py          |  8 ++++----
 src/palimpzest/query/operators/logical.py              | 10 ++--------
 3 files changed, 8 insertions(+), 14 deletions(-)

diff --git a/src/palimpzest/query/generators/generators.py b/src/palimpzest/query/generators/generators.py
index e7d8e1ac..9ad5c9b1 100644
--- a/src/palimpzest/query/generators/generators.py
+++ b/src/palimpzest/query/generators/generators.py
@@ -12,7 +12,7 @@
 from abc import ABC, abstractmethod
 from collections import Counter
 from string import Formatter
-from typing import Any, Generic, Tuple, TypeVar
+from typing import Any, Generic, TypeVar
 
 from colorama import Fore, Style
 from openai import OpenAI
@@ -40,7 +40,7 @@
 from palimpzest.utils.sandbox import API
 
 # DEFINITIONS
-GenerationOutput = Tuple[dict, str | None, GenerationStats]
+GenerationOutput = tuple[dict, str | None, GenerationStats]
 ContextType = TypeVar("ContextType")
 InputType = TypeVar("InputType")
 
diff --git a/src/palimpzest/query/operators/code_synthesis_convert.py b/src/palimpzest/query/operators/code_synthesis_convert.py
index 5e9d610f..650fc613 100644
--- a/src/palimpzest/query/operators/code_synthesis_convert.py
+++ b/src/palimpzest/query/operators/code_synthesis_convert.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Dict, Tuple
+from typing import Any
 
 from palimpzest.constants import Cardinality, GPT_4o_MODEL_CARD, Model
 from palimpzest.core.data.dataclasses import GenerationStats, OperatorCostEstimates
@@ -15,9 +15,9 @@
 FieldName = str
 CodeName = str
 Code = str
-DataRecordDict = Dict[str, Any]
-Exemplar = Tuple[DataRecordDict, DataRecordDict]
-CodeEnsemble = Dict[CodeName, Code]
+DataRecordDict = dict[str, Any]
+Exemplar = tuple[DataRecordDict, DataRecordDict]
+CodeEnsemble = dict[CodeName, Code]
 
 
 class CodeSynthesisConvert(LLMConvert):
diff --git a/src/palimpzest/query/operators/logical.py b/src/palimpzest/query/operators/logical.py
index 28f50c80..4bb739de 100644
--- a/src/palimpzest/query/operators/logical.py
+++ b/src/palimpzest/query/operators/logical.py
@@ -137,14 +137,8 @@ def get_logical_op_params(self) -> dict:
 class BaseScan(LogicalOperator):
     """A BaseScan is a logical operator that represents a scan of a particular data source."""
 
-    def __init__(self, dataset_id: str, *args, **kwargs):
-        if kwargs.get("input_schema") is not None:
-            raise Exception(
-                f"BaseScan must be initialized with `input_schema=None` but was initialized with "
-                f"`input_schema={kwargs.get('input_schema')}`"
-            )
-
-        super().__init__(*args, **kwargs)
+    def __init__(self, dataset_id: str, output_schema: Schema):
+        super().__init__(output_schema=output_schema)
         self.dataset_id = dataset_id
 
     def __str__(self):