Redesign BaseModel with unified inheritance and safety validation (#195)

- Replace separate _BaseModelPyodide/_BaseModelStd with single BaseModel class
- Use __init_subclass__ for consistent dataclass transformation across environments
- Add comprehensive inheritance validation:
- Prevent dataclass mixins while inheriting from BaseModel
- Block dangerous auto_dataclass=False when parent was auto-dataclassed
- Require BaseModel or BaseModel subclass as primary base
- Support dataclass options (frozen=True, etc.) via kwargs passthrough
- Add __auto_dataclass tracking attribute for inheritance state
- Preserve custom __init__ methods while maintaining dataclass benefits

Fix custom dataclass coder support for list[Output] returns:
- Auto-register DataclassCoder when BaseModel subclasses are created
- Add custom_output test runner and schema validation
- Fix Go test expected output formatting with JSON round-trip serialization
- Remove manual dataclass_coder import requirement for users

Remove obsolete pyodide-specific code and add comprehensive test suite
covering all inheritance scenarios. This creates a safer, more predictable
BaseModel system that prevents broken field inheritance while enabling
seamless custom dataclass serialization.

Author: tempusfrangit

internal/tests/coder_test.go

@@ -88,3 +88,45 @@ func TestPredictionChatCoderSucceeded(t *testing.T) {
 	assert.Equal(t, expectedOutput, predictionResponse.Output)
 	assert.Equal(t, server.PredictionSucceeded, predictionResponse.Status)
 }
+
+func TestPredictionCustomOutputCoder(t *testing.T) {
+	t.Parallel()
+	if *legacyCog {
+		t.Skip("legacy Cog does not support custom coder")
+	}
+
+	runtimeServer := setupCogRuntime(t, cogRuntimeServerConfig{
+		procedureMode:    false,
+		explicitShutdown: true,
+		uploadURL:        "",
+		module:           "custom_output",
+		predictorClass:   "Predictor",
+	})
+	waitForSetupComplete(t, runtimeServer, server.StatusReady, server.SetupSucceeded)
+
+	input := map[string]any{"i": 3}
+	req := httpPredictionRequest(t, runtimeServer, server.PredictionRequest{Input: input})
+	resp, err := http.DefaultClient.Do(req)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+	assert.Equal(t, http.StatusOK, resp.StatusCode)
+	body, err := io.ReadAll(resp.Body)
+	require.NoError(t, err)
+	var predictionResponse server.PredictionResponse
+	err = json.Unmarshal(body, &predictionResponse)
+	require.NoError(t, err)
+
+	// Create expected output using JSON round-trip to match server serialization
+	expectedItems := []map[string]any{
+		{"x": 3, "y": "a"},
+		{"x": 2, "y": "a"},
+		{"x": 1, "y": "a"},
+	}
+	expectedJSON, err := json.Marshal(expectedItems)
+	require.NoError(t, err)
+	var expectedOutput []any
+	err = json.Unmarshal(expectedJSON, &expectedOutput)
+	require.NoError(t, err)
+	assert.Equal(t, expectedOutput, predictionResponse.Output)
+	assert.Equal(t, server.PredictionSucceeded, predictionResponse.Status)
+}

python/coglet/api.py

@@ -1,5 +1,4 @@
 import pathlib
-import sys
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, is_dataclass
 from typing import Any, AsyncIterator, Iterator, List, Optional, Type, TypeVar, Union
@@ -97,39 +96,48 @@ class Input:
     deprecated: Optional[bool] = None
 
 
-# pyodide does not recognise `BaseModel` with the `__new__` keyword as a dataclass while regular python does.
-# to get around this, we hijack `__init__subclass` instead to make sure the subclass of a base model is recognised
-# as a dataclass. In addition to this, we provide this `BaseModel` only on pyodide instances, normal python still gets
-# the regular `BaseModel``.
-class _BaseModelPyodide:
-    def __init_subclass__(cls, **kwargs):
-        dc_keys = {
-            'init',
-            'repr',
-            'eq',
-            'order',
-            'unsafe_hash',
-            'frozen',
-            'match_args',
-            'kw_only',
-            'slots',
-            'weakref_slot',
-        }
-        dc_opts = {k: kwargs.pop(k) for k in list(kwargs) if k in dc_keys}
-        super().__init_subclass__(**kwargs)
-        if not is_dataclass(cls):
-            dataclass(**dc_opts)(cls)
-
-
-class _BaseModelStd:
-    def __new__(cls, *args, **kwargs):
-        # This does not work with frozen=True
-        # Also user might want to mutate the output class
-        dcls = dataclass()(cls)
-        return super().__new__(dcls)
-
-
-BaseModel = _BaseModelPyodide if 'pyodide' in sys.modules else _BaseModelStd
+class BaseModel:
+    def __init_subclass__(
+        cls, *, auto_dataclass: bool = True, init: bool = True, **kwargs
+    ):
+        # BaseModel is parented to `object` so we have nothing to pass up to it, we pass the kwargs to dataclass() only.
+        super().__init_subclass__()
+
+        # For sanity, the primary base class must inherit from BaseModel
+        if not issubclass(cls.__bases__[0], BaseModel):
+            raise TypeError(
+                f'Primary base class of "{cls.__name__}" must inherit from BaseModel'
+            )
+        elif not auto_dataclass:
+            try:
+                if (
+                    cls.__bases__[0] != BaseModel
+                    and cls.__bases__[0].__auto_dataclass is True  # type: ignore[attr-defined]
+                ):
+                    raise ValueError(
+                        f'Primary base class of "{cls.__name__}" ("{cls.__bases__[0].__name__}") has auto_dataclass=True, but "{cls.__name__}" has auto_dataclass=False. This creates broken field inheritance.'
+                    )
+            except AttributeError:
+                raise RuntimeError(
+                    f'Primary base class of "{cls.__name__}" is a child of a child of `BaseModel`, but `auto_dataclass` tracking does not exist. This is likely a bug or other programming error.'
+                )
+
+        for base in cls.__bases__[1:]:
+            if is_dataclass(base):
+                raise TypeError(
+                    f'Cannot mixin dataclass "{base.__name__}" while inheriting from `BaseModel`'
+                )
+
+        # Once manual dataclass handling is enabled, we never apply the auto dataclass logic again,
+        # it becomes the responsibility of the user to ensure that all dataclass semantics are handled.
+        if not auto_dataclass:
+            cls.__auto_dataclass = False  # type: ignore[attr-defined]
+            return
+
+        # all children should be dataclass'd, this is the only way to ensure that the dataclass inheritence
+        # is handled properly.
+        dataclass(init=init, **kwargs)(cls)
+        cls.__auto_dataclass = True  # type: ignore[attr-defined]
 
 
 ########################################

python/tests/runners/custom_output.py

@@ -0,0 +1,17 @@
+from cog import BaseModel, BasePredictor
+
+
+class CustomOut(BaseModel):
+    x: int
+    y: str
+
+
+class Predictor(BasePredictor):
+    test_inputs = {'i': 3}
+
+    def predict(self, i: int) -> list[CustomOut]:
+        outputs: list[CustomOut] = []
+        while i > 0:
+            outputs.append(CustomOut(x=i, y='a'))
+            i -= 1
+        return outputs