Remove EnumParameterSchema

tool2schema/parameter_schema.py

@@ -82,6 +82,14 @@ def _get_default(self) -> Any:
         # Parameter.empty to indicate that the default value is not present
         return Parameter.empty
 
+    def add_enum(self, values: list) -> None:
+        """
+        Convert this parameter to an enumeration type.
+
+        :param values: List of unique enumeration values.
+        """
+        self.type_schema = EnumTypeSchema(values)
+
     def to_json(self) -> dict:
         """
         Return the json schema for this parameter.
@@ -95,14 +103,3 @@ def to_json(self) -> dict:
         json = {f: v for f, v in fields.items() if v != Parameter.empty}
 
         return json
-
-
-class EnumParameterSchema(ParameterSchema):
-    """
-    Parameter schema for enumeration types manually added via add_enum.
-    """
-
-    def __init__(
-        self, values: list, parameter: Parameter, index: int, config: Config, docstring: str = None
-    ):
-        super().__init__(EnumTypeSchema(values), parameter, index, config, docstring)

tool2schema/schema.py

@@ -9,7 +9,7 @@
 
 import tool2schema
 from tool2schema.config import Config
-from tool2schema.parameter_schema import EnumParameterSchema, ParameterSchema
+from tool2schema.parameter_schema import ParameterSchema
 
 
 class SchemaType(Enum):
@@ -150,10 +150,7 @@ def add_enum(self, n: str, enum: list) -> "FunctionSchema":
         :param n: The name of the parameter with the enum values
         :param enum: The list of values for the enum parameter
         """
-        p = self._all_parameter_schemas[n]
-        self._all_parameter_schemas[n] = EnumParameterSchema(
-            enum, p.parameter, p.index, self.config, p.docstring
-        )
+        self._all_parameter_schemas[n].add_enum(enum)
         return self
 
     def _get_schema(self) -> dict: