Added support for frozen_default for dataclass_transform.

msfterictraut · msfterictraut · commit bedc124c394f · 2022-12-02T04:35:05.000-08:00
diff --git a/packages/pyright-internal/src/analyzer/dataClasses.ts b/packages/pyright-internal/src/analyzer/dataClasses.ts
@@ -713,6 +713,7 @@ export function validateDataClassTransformDecorator(
         keywordOnlyParams: false,
         generateEq: true,
         generateOrder: false,
+        frozen: false,
         fieldDescriptorNames: [],
     };
 
@@ -780,6 +781,24 @@ export function validateDataClassTransformDecorator(
                 break;
             }
 
+            case 'frozen_default': {
+                const value = evaluateStaticBoolExpression(
+                    arg.valueExpression,
+                    fileInfo.executionEnvironment,
+                    fileInfo.definedConstants
+                );
+                if (value === undefined) {
+                    evaluator.addError(
+                        Localizer.Diagnostic.dataClassTransformExpectedBoolLiteral(),
+                        arg.valueExpression
+                    );
+                    return;
+                }
+
+                behaviors.frozen = value;
+                break;
+            }
+
             // Earlier versions of the dataclass_transform spec used the name "field_descriptors"
             // rather than "field_specifiers". The older name is now deprecated but still supported
             // for the time being because some libraries shipped with the older __dataclass_transform__
@@ -858,6 +877,7 @@ export function getDataclassDecoratorBehaviors(type: Type): DataClassBehaviors |
             keywordOnlyParams: false,
             generateEq: true,
             generateOrder: false,
+            frozen: false,
             fieldDescriptorNames: ['dataclasses.field', 'dataclasses.Field'],
         };
     }
@@ -998,7 +1018,8 @@ export function applyDataClassClassBehaviorOverrides(
     evaluator: TypeEvaluator,
     errorNode: ParseNode,
     classType: ClassType,
-    args: FunctionArgument[]
+    args: FunctionArgument[],
+    defaultBehaviors: DataClassBehaviors
 ) {
     let sawFrozenArg = false;
 
@@ -1015,7 +1036,7 @@ export function applyDataClassClassBehaviorOverrides(
     // If there was no frozen argument, it is implicitly false. This will
     // validate that we're not overriding a frozen class with a non-frozen class.
     if (!sawFrozenArg) {
-        applyDataClassBehaviorOverrideValue(evaluator, errorNode, classType, 'frozen', false);
+        applyDataClassBehaviorOverrideValue(evaluator, errorNode, classType, 'frozen', defaultBehaviors.frozen);
     }
 }
 
@@ -1034,6 +1055,10 @@ export function applyDataClassDefaultBehaviors(classType: ClassType, defaultBeha
     if (defaultBehaviors.generateOrder) {
         classType.details.flags |= ClassTypeFlags.SynthesizedDataClassOrder;
     }
+
+    if (defaultBehaviors.frozen) {
+        classType.details.flags |= ClassTypeFlags.FrozenDataClass;
+    }
 }
 
 export function applyDataClassDecorator(
@@ -1045,7 +1070,5 @@ export function applyDataClassDecorator(
 ) {
     applyDataClassDefaultBehaviors(classType, defaultBehaviors);
 
-    if (callNode?.arguments) {
-        applyDataClassClassBehaviorOverrides(evaluator, errorNode, classType, callNode.arguments);
-    }
+    applyDataClassClassBehaviorOverrides(evaluator, errorNode, classType, callNode?.arguments ?? [], defaultBehaviors);
 }
diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts
@@ -15576,7 +15576,13 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
 
         if (dataClassBehaviors) {
             applyDataClassDefaultBehaviors(classType, dataClassBehaviors);
-            applyDataClassClassBehaviorOverrides(evaluatorInterface, node.name, classType, initSubclassArgs);
+            applyDataClassClassBehaviorOverrides(
+                evaluatorInterface,
+                node.name,
+                classType,
+                initSubclassArgs,
+                dataClassBehaviors
+            );
         }
 
         // Run any class hooks that depend on this class.
diff --git a/packages/pyright-internal/src/analyzer/types.ts b/packages/pyright-internal/src/analyzer/types.ts
@@ -461,6 +461,7 @@ export interface DataClassBehaviors {
     keywordOnlyParams: boolean;
     generateEq: boolean;
     generateOrder: boolean;
+    frozen: boolean;
     fieldDescriptorNames: string[];
 }
 
diff --git a/packages/pyright-internal/src/tests/samples/dataclassTransform1.py b/packages/pyright-internal/src/tests/samples/dataclassTransform1.py
@@ -69,3 +69,24 @@ class Customer2Subclass(Customer2, frozen=True):
 c2_2 = Customer2(0, "John")
 
 v2 = c2_1 < c2_2
+
+
+@dataclass_transform(kw_only_default=True, order_default=True, frozen_default=True)
+def create_model_frozen(cls: _T) -> _T:
+    ...
+
+@create_model_frozen
+class Customer3:
+    id: int
+    name: str
+
+# This should generate an error because a non-frozen class
+# cannot inherit from a frozen class.
+@create_model
+class Customer3Subclass(Customer3):
+    age: int
+
+c3_1 = Customer3(id=2, name="hi")
+
+# This should generate an error because Customer3 is frozen.
+c3_1.id = 4
diff --git a/packages/pyright-internal/src/tests/samples/dataclassTransform2.py b/packages/pyright-internal/src/tests/samples/dataclassTransform2.py
@@ -80,3 +80,22 @@ class Customer2(ModelBase, order=True):
 # This should generate an error because Customer2 supports
 # keyword-only parameters for its constructor.
 c2_3 = Customer2(0, "John")
+
+
+
+@dataclass_transform(frozen_default=True)
+class ModelMetaFrozen(type):
+    pass
+
+class ModelBaseFrozen(metaclass=ModelMetaFrozen):
+    ...
+
+class Customer3(ModelBaseFrozen):
+    id: int
+    name: str
+
+
+c3_1 = Customer3(id=2, name="hi")
+
+# This should generate an error because Customer3 is frozen.
+c3_1.id = 4
diff --git a/packages/pyright-internal/src/tests/samples/dataclassTransform3.py b/packages/pyright-internal/src/tests/samples/dataclassTransform3.py
@@ -11,6 +11,7 @@ def __dataclass_transform__(
     eq_default: bool = True,
     order_default: bool = False,
     kw_only_default: bool = False,
+    frozen_default: bool = False,
     field_specifiers: Tuple[Union[type, Callable[..., Any]], ...] = (()),
 ) -> Callable[[_T], _T]:
     return lambda a: a
@@ -108,4 +109,18 @@ def __init_subclass__(
 class GenericCustomer(GenericModelBase[int]):
     id: int = model_field()
 
-gc_1 = GenericCustomer(id=3)
+gc_1 = GenericCustomer(id=3)
+
+@__dataclass_transform__(frozen_default=True)
+class ModelBaseFrozen:
+    not_a_field: str
+
+class Customer3(ModelBaseFrozen):
+    id: int
+    name: str
+
+
+c3_1 = Customer3(id=2, name="hi")
+
+# This should generate an error because Customer3 is frozen.
+c3_1.id = 4
diff --git a/packages/pyright-internal/src/tests/typeEvaluator3.test.ts b/packages/pyright-internal/src/tests/typeEvaluator3.test.ts
@@ -1383,19 +1383,19 @@ test('Decorator7', () => {
 test('DataclassTransform1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['dataclassTransform1.py']);
 
-    TestUtils.validateResults(analysisResults, 4);
+    TestUtils.validateResults(analysisResults, 6);
 });
 
 test('DataclassTransform2', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['dataclassTransform2.py']);
 
-    TestUtils.validateResults(analysisResults, 5);
+    TestUtils.validateResults(analysisResults, 6);
 });
 
 test('DataclassTransform3', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['dataclassTransform3.py']);
 
-    TestUtils.validateResults(analysisResults, 5);
+    TestUtils.validateResults(analysisResults, 6);
 });
 
 test('DataclassTransform4', () => {
diff --git a/packages/pyright-internal/typeshed-fallback/stdlib/typing_extensions.pyi b/packages/pyright-internal/typeshed-fallback/stdlib/typing_extensions.pyi
@@ -227,6 +227,7 @@ else:
         eq_default: bool = ...,
         order_default: bool = ...,
         kw_only_default: bool = ...,
+        frozen_default: bool = ...,
         field_specifiers: tuple[type[Any] | Callable[..., Any], ...] = ...,
         **kwargs: object,
     ) -> IdentityFunction: ...

Original file line number	Diff line number	Diff line change
`@@ -461,6 +461,7 @@ export interface DataClassBehaviors {`
`461`	`461`	`keywordOnlyParams: boolean;`
`462`	`462`	`generateEq: boolean;`
`463`	`463`	`generateOrder: boolean;`
	`464`	`+ frozen: boolean;`
`464`	`465`	`fieldDescriptorNames: string[];`
`465`	`466`	`}`
`466`	`467`