import { assert } from '../common/debug' ; import { DiagnosticRule } from '../common/diagnosticRules' ; import { Localizer } from '../localization/localize' ; import { ArgumentCategory , CallNode , ClassNode , ExpressionNode , NameNode , ParameterCategory , ParseNode , ParseNodeType , TypeAnnotationNode , } from '../parser/parseNodes' ; import * as AnalyzerNodeInfo from './analyzerNodeInfo' ; import { DeclarationType } from './declaration' ; import { updateNamedTupleBaseClass } from './namedTuples' ; import { evaluateStaticBoolExpression } from './staticExpressions' ; import { Symbol , SymbolFlags } from './symbol' ; import { EvaluatorFlags , FunctionArgument , TypeEvaluator } from './typeEvaluatorTypes' ; import { AnyType , ClassType , ClassTypeFlags , DataClassBehaviors , DataClassEntry , FunctionParameter , FunctionType , FunctionTypeFlags , isClass , isClassInstance , isFunction , isInstantiableClass , isOverloadedFunction , NoneType , Type , UnknownType , } from './types' ; import { applySolvedTypeVars , buildTypeVarMapFromSpecializedClass , convertToInstance , isLiteralType , specializeTupleClass , synthesizeTypeVarForSelfCls , } from './typeUtils' ; export function synthesizeDataClassMethods ( evaluator , node , classType , skipSynthesizeInit , skipSynthesizeHash ) { assert ( ClassType . isDataClass ( classType ) ) ; const classTypeVar = synthesizeTypeVarForSelfCls ( classType , true ) ; const newType = FunctionType . createInstance ( '__new__' , '' , '' , FunctionTypeFlags . ConstructorMethod | FunctionTypeFlags . SynthesizedMethod ) ; const initType = FunctionType . createInstance ( '__init__' , '' , '' , FunctionTypeFlags . SynthesizedMethod ) ; FunctionType . addParameter ( newType , { category : ParameterCategory . Simple , name : 'cls' , type : classTypeVar , hasDeclaredType : true , } ) ; FunctionType . addDefaultParameters ( newType ) ; newType . details . declaredReturnType = convertToInstance ( classTypeVar ) ; const selfParam = { category : ParameterCategory . Simple , name : 'self' , type : synthesizeTypeVarForSelfCls ( classType , false ) , hasDeclaredType : true , } ; FunctionType . addParameter ( initType , selfParam ) ; initType . details . declaredReturnType = NoneType . createInstance ( ) ; const localDataClassEntries = [ ] ; const fullDataClassEntries = [ ] ; const allAncestorsKnown = addInheritedDataClassEntries ( classType , fullDataClassEntries ) ; if ( ! allAncestorsKnown ) { FunctionType . addDefaultParameters ( initType ) ; } type TypeEvaluator = ( ) => Type ; const localEntryTypeEvaluator = [ ] ; let sawKeywordOnlySeparator = false ; node . suite . statements . forEach ( ( statementList ) => { if ( statementList . nodeType === ParseNodeType . StatementList ) { statementList . statements . forEach ( ( statement ) => { let variableNameNode ; let aliasName ; let variableTypeEvaluator ; let hasDefaultValue = false ; let isKeywordOnly = ClassType . isDataClassKeywordOnlyParams ( classType ) || sawKeywordOnlySeparator ; let defaultValueExpression ; let includeInInit = true ; if ( statement . nodeType === ParseNodeType . Assignment ) { if ( statement . leftExpression . nodeType === ParseNodeType . TypeAnnotation && statement . leftExpression . valueExpression . nodeType === ParseNodeType . Name ) { variableNameNode = statement . leftExpression . valueExpression ; variableTypeEvaluator = ( ) => evaluator . getTypeOfAnnotation ( ( statement . leftExpression as TypeAnnotationNode ) . typeAnnotation , { isVariableAnnotation : true , allowFinal : true , allowClassVar : true , } ) ; } hasDefaultValue = true ; defaultValueExpression = statement . rightExpression ; if ( statement . rightExpression . nodeType === ParseNodeType . Call ) { const callType = evaluator . getTypeOfExpression ( statement . rightExpression . leftExpression , undefined , EvaluatorFlags . DoNotSpecialize ) . type ; if ( isDataclassFieldConstructor ( callType , classType . details . dataClassBehaviors ?. fieldDescriptorNames || [ ] ) ) { const initArg = statement . rightExpression . arguments . find ( ( arg ) => arg . name ?. value === 'init' ) ; if ( initArg && initArg . valueExpression ) { const value = evaluateStaticBoolExpression ( initArg . valueExpression , AnalyzerNodeInfo . getFileInfo ( node ) . executionEnvironment ) ; if ( value === false ) { includeInInit = false ; } } else { let callTarget ; if ( isFunction ( callType ) ) { callTarget = callType ; } else if ( isOverloadedFunction ( callType ) ) { callTarget = evaluator . getBestOverloadForArguments ( statement . rightExpression , callType , statement . rightExpression . arguments ) ; } else if ( isInstantiableClass ( callType ) ) { const initCall = evaluator . getBoundMethod ( callType , '__init__' ) ; if ( initCall ) { if ( isFunction ( initCall ) ) { callTarget = initCall ; } else if ( isOverloadedFunction ( initCall ) ) { callTarget = evaluator . getBestOverloadForArguments ( statement . rightExpression , initCall , statement . rightExpression . arguments ) ; } } } if ( callTarget ) { const initParam = callTarget . details . parameters . find ( ( p ) => p . name === 'init' ) ; if ( initParam && initParam . defaultValueExpression && initParam . hasDeclaredType ) { if ( isClass ( initParam . type ) && ClassType . isBuiltIn ( initParam . type , 'bool' ) && isLiteralType ( initParam . type ) ) { if ( initParam . type . literalValue === false ) { includeInInit = false ; } } } } } const kwOnlyArg = statement . rightExpression . arguments . find ( ( arg ) => arg . name ?. value === 'kw_only' ) ; if ( kwOnlyArg && kwOnlyArg . valueExpression ) { const value = evaluateStaticBoolExpression ( kwOnlyArg . valueExpression , AnalyzerNodeInfo . getFileInfo ( node ) . executionEnvironment ) ; if ( value === false ) { isKeywordOnly = false ; } else if ( value === true ) { isKeywordOnly = true ; } } hasDefaultValue = statement . rightExpression . arguments . some ( ( arg ) => arg . name ?. value === 'default' || arg . name ?. value === 'default_factory' || arg . name ?. value === 'factory' ) ; const aliasArg = statement . rightExpression . arguments . find ( ( arg ) => arg . name ?. value === 'alias' ) ; if ( aliasArg ) { const valueType = evaluator . getTypeOfExpression ( aliasArg . valueExpression ) . type ; if ( isClassInstance ( valueType ) && ClassType . isBuiltIn ( valueType , 'str' ) && isLiteralType ( valueType ) ) { aliasName = valueType . literalValue as string ; } } } } } else if ( statement . nodeType === ParseNodeType . TypeAnnotation ) { if ( statement . valueExpression . nodeType === ParseNodeType . Name ) { variableNameNode = statement . valueExpression ; variableTypeEvaluator = ( ) => evaluator . getTypeOfAnnotation ( statement . typeAnnotation , { isVariableAnnotation : true , allowFinal : true , allowClassVar : true , } ) ; if ( statement . valueExpression . value === '_' ) { const annotatedType = variableTypeEvaluator ( ) ; if ( isClassInstance ( annotatedType ) && ClassType . isBuiltIn ( annotatedType , 'KW_ONLY' ) ) { sawKeywordOnlySeparator = true ; variableNameNode = undefined ; variableTypeEvaluator = undefined ; } } } } if ( variableNameNode && variableTypeEvaluator ) { const variableName = variableNameNode . value ; const variableSymbol = classType . details . fields . get ( variableName ) ; const isFinal = variableSymbol ?. getDeclarations ( ) . some ( ( decl ) => decl . type === DeclarationType . Variable && decl . isFinal ) ; if ( variableSymbol ?. isClassVar ( ) && ! isFinal ) { const index = fullDataClassEntries . findIndex ( ( p ) => p . name === variableName ) ; if ( index >= 0 ) { fullDataClassEntries . splice ( index , 1 ) ; } const dataClassEntry = { name : variableName , alias : aliasName , isKeywordOnly : false , hasDefault : hasDefaultValue , defaultValueExpression , includeInInit , type : UnknownType . create ( ) , isClassVar : true , } ; localDataClassEntries . push ( dataClassEntry ) ; } else { const dataClassEntry = { name : variableName , alias : aliasName , isKeywordOnly , hasDefault : hasDefaultValue , defaultValueExpression , includeInInit , type : UnknownType . create ( ) , isClassVar : false , } ; localEntryTypeEvaluator . push ( { entry : dataClassEntry , evaluator : variableTypeEvaluator } ) ; let insertIndex = localDataClassEntries . findIndex ( ( e ) => e . name === variableName ) ; if ( insertIndex >= 0 ) { localDataClassEntries [ insertIndex ] = dataClassEntry ; } else { localDataClassEntries . push ( dataClassEntry ) ; } insertIndex = fullDataClassEntries . findIndex ( ( p ) => p . name === variableName ) ; if ( insertIndex >= 0 ) { fullDataClassEntries [ insertIndex ] = dataClassEntry ; } else { fullDataClassEntries . push ( dataClassEntry ) ; insertIndex = fullDataClassEntries . length - 1 ; } if ( ! isKeywordOnly && includeInInit && ! skipSynthesizeInit && ! hasDefaultValue ) { const firstDefaultValueIndex = fullDataClassEntries . findIndex ( ( p ) => p . hasDefault && p . includeInInit && ! p . isKeywordOnly ) ; if ( firstDefaultValueIndex >= 0 && firstDefaultValueIndex < insertIndex ) { evaluator . addError ( Localizer . Diagnostic . dataClassFieldWithDefault ( ) , variableNameNode ) ; } } } } } ) ; } } ) ; classType . details . dataClassEntries = localDataClassEntries ; localEntryTypeEvaluator . forEach ( ( entryEvaluator ) => { entryEvaluator . entry . type = entryEvaluator . evaluator ( ) ; } ) ; const symbolTable = classType . details . fields ; const keywordOnlyParams = [ ] ; if ( ! skipSynthesizeInit && allAncestorsKnown ) { fullDataClassEntries . forEach ( ( entry ) => { if ( entry . includeInInit ) { const functionParam = { category : ParameterCategory . Simple , name : entry . alias || entry . name , hasDefault : entry . hasDefault , defaultValueExpression : entry . defaultValueExpression , type : entry . type , hasDeclaredType : true , } ; if ( entry . isKeywordOnly ) { keywordOnlyParams . push ( functionParam ) ; } else { FunctionType . addParameter ( initType , functionParam ) ; } } } ) ; if ( keywordOnlyParams . length > 0 ) { FunctionType . addParameter ( initType , { category : ParameterCategory . VarArgList , type : AnyType . create ( ) , } ) ; keywordOnlyParams . forEach ( ( param ) => { FunctionType . addParameter ( initType , param ) ; } ) ; } symbolTable . set ( '__init__' , Symbol . createWithType ( SymbolFlags . ClassMember , initType ) ) ; symbolTable . set ( '__new__' , Symbol . createWithType ( SymbolFlags . ClassMember , newType ) ) ; } const strType = evaluator . getBuiltInType ( node , 'str' ) ; const tupleClassType = evaluator . getBuiltInType ( node , 'tuple' ) ; if ( tupleClassType && isInstantiableClass ( tupleClassType ) && strType && isInstantiableClass ( strType ) && ! symbolTable . has ( '__match_args__' ) ) { const matchArgsNames = [ ] ; fullDataClassEntries . forEach ( ( entry ) => { if ( entry . includeInInit && ! entry . isKeywordOnly ) { matchArgsNames . push ( entry . name ) ; } } ) ; const literalTypes = matchArgsNames . map ( ( name ) => { return ClassType . cloneAsInstance ( ClassType . cloneWithLiteral ( strType , name ) ) ; } ) ; const matchArgsType = ClassType . cloneAsInstance ( specializeTupleClass ( tupleClassType , literalTypes ) ) ; symbolTable . set ( '__match_args__' , Symbol . createWithType ( SymbolFlags . ClassMember , matchArgsType ) ) ; } const synthesizeComparisonMethod = ( operator , paramType ) => { const operatorMethod = FunctionType . createInstance ( operator , '' , '' , FunctionTypeFlags . SynthesizedMethod ) ; FunctionType . addParameter ( operatorMethod , selfParam ) ; FunctionType . addParameter ( operatorMethod , { category : ParameterCategory . Simple , name : 'x' , type : paramType , hasDeclaredType : true , } ) ; operatorMethod . details . declaredReturnType = evaluator . getBuiltInObject ( node , 'bool' ) ; symbolTable . set ( operator , Symbol . createWithType ( SymbolFlags . ClassMember , operatorMethod ) ) ; } ; if ( ! ClassType . isSkipSynthesizedDataClassEq ( classType ) ) { synthesizeComparisonMethod ( '__eq__' , evaluator . getBuiltInObject ( node , 'object' ) ) ; } if ( ClassType . isSynthesizedDataclassOrder ( classType ) ) { const objType = ClassType . cloneAsInstance ( classType ) ; [ '__lt__' , '__le__' , '__gt__' , '__ge__' ] . forEach ( ( operator ) => { synthesizeComparisonMethod ( operator , objType ) ; } ) ; } let synthesizeHashFunction = ! ClassType . isSkipSynthesizedDataClassEq ( classType ) && ClassType . isFrozenDataClass ( classType ) ; const synthesizeHashNone = ! ClassType . isSkipSynthesizedDataClassEq ( classType ) && ! ClassType . isFrozenDataClass ( classType ) ; if ( skipSynthesizeHash ) { synthesizeHashFunction = false ; } if ( ClassType . isSynthesizeDataClassUnsafeHash ( classType ) ) { synthesizeHashFunction = true ; } if ( synthesizeHashFunction ) { const hashMethod = FunctionType . createInstance ( '__hash__' , '' , '' , FunctionTypeFlags . SynthesizedMethod ) ; FunctionType . addParameter ( hashMethod , selfParam ) ; hashMethod . details . declaredReturnType = evaluator . getBuiltInObject ( node , 'int' ) ; symbolTable . set ( '__hash__' , Symbol . createWithType ( SymbolFlags . ClassMember , hashMethod ) ) ; } else if ( synthesizeHashNone && ! skipSynthesizeHash ) { symbolTable . set ( '__hash__' , Symbol . createWithType ( SymbolFlags . ClassMember , NoneType . createInstance ( ) ) ) ; } let dictType = evaluator . getBuiltInType ( node , 'dict' ) ; if ( isInstantiableClass ( dictType ) ) { dictType = ClassType . cloneAsInstance ( ClassType . cloneForSpecialization ( dictType , [ evaluator . getBuiltInObject ( node , 'str' ) , AnyType . create ( ) ] , true ) ) ; } symbolTable . set ( '__dataclass_fields__' , Symbol . createWithType ( SymbolFlags . ClassMember , dictType ) ) ; if ( ClassType . isGeneratedDataClassSlots ( classType ) && classType . details . localSlotsNames === undefined ) { classType . details . localSlotsNames = localDataClassEntries . map ( ( entry ) => entry . name ) ; } updateNamedTupleBaseClass ( classType , fullDataClassEntries . map ( ( entry ) => entry . type ) , true ) ; } function addInheritedDataClassEntries ( classType , entries ) { let allAncestorsAreKnown = true ; for ( let i = classType . details . mro . length - 1 ; i >= 0 ; i -- ) { const mroClass = classType . details . mro [ i ] ; if ( isInstantiableClass ( mroClass ) ) { const typeVarMap = buildTypeVarMapFromSpecializedClass ( mroClass , false ) ; const dataClassEntries = ClassType . getDataClassEntries ( mroClass ) ; dataClassEntries . forEach ( ( entry ) => { const existingIndex = entries . findIndex ( ( e ) => e . name === entry . name ) ; const updatedEntry = { ... entry } ; updatedEntry . type = applySolvedTypeVars ( updatedEntry . type , typeVarMap ) ; if ( entry . isClassVar ) { if ( existingIndex >= 0 ) { entries . splice ( existingIndex , 1 ) ; } } else if ( existingIndex >= 0 ) { entries [ existingIndex ] = updatedEntry ; } else { entries . push ( updatedEntry ) ; } } ) ; } else { allAncestorsAreKnown = false ; } } return allAncestorsAreKnown ; } function isDataclassFieldConstructor ( type , fieldDescriptorNames ) { let callName ; if ( isFunction ( type ) ) { callName = type . details . fullName ; } else if ( isOverloadedFunction ( type ) ) { callName = type . overloads [ 0 ] . details . fullName ; } else if ( isInstantiableClass ( type ) ) { callName = type . details . fullName ; } if ( ! callName ) { return false ; } return fieldDescriptorNames . some ( ( name ) => name === callName ) ; } export function validateDataClassTransformDecorator ( evaluator , node ) { if ( node . nodeType !== ParseNodeType . Call ) { return undefined ; } const behaviors = { keywordOnlyParams : false , generateEq : true , generateOrder : false , fieldDescriptorNames : [ ] , } ; const fileInfo = AnalyzerNodeInfo . getFileInfo ( node ) ; node . arguments . forEach ( ( arg ) => { if ( ! arg . name ) { return ; } if ( arg . argumentCategory !== ArgumentCategory . Simple ) { return ; } switch ( arg . name . value ) { case 'kw_only_default' : { const value = evaluateStaticBoolExpression ( arg . valueExpression , fileInfo . executionEnvironment ) ; if ( value === undefined ) { return ; } behaviors . keywordOnlyParams = value ; break ; } case 'eq_default' : { const value = evaluateStaticBoolExpression ( arg . valueExpression , fileInfo . executionEnvironment ) ; if ( value === undefined ) { return ; } behaviors . generateEq = value ; break ; } case 'order_default' : { const value = evaluateStaticBoolExpression ( arg . valueExpression , fileInfo . executionEnvironment ) ; if ( value === undefined ) { return ; } behaviors . generateOrder = value ; break ; } case 'field_descriptors' : { const valueType = evaluator . getTypeOfExpression ( arg . valueExpression ) . type ; if ( ! isClassInstance ( valueType ) || ! ClassType . isBuiltIn ( valueType , 'tuple' ) || ! valueType . tupleTypeArguments || valueType . tupleTypeArguments . some ( ( entry ) => ! isInstantiableClass ( entry ) && ! isFunction ( entry ) && ! isOverloadedFunction ( entry ) ) ) { return ; } if ( ! behaviors . fieldDescriptorNames ) { behaviors . fieldDescriptorNames = [ ] ; } valueType . tupleTypeArguments . forEach ( ( arg ) => { if ( isInstantiableClass ( arg ) || isFunction ( arg ) ) { behaviors . fieldDescriptorNames . push ( arg . details . fullName ) ; } else if ( isOverloadedFunction ( arg ) ) { behaviors . fieldDescriptorNames . push ( arg . overloads [ 0 ] . details . fullName ) ; } } ) ; break ; } default : break ; } } ) ; return behaviors ; } export function getDataclassDecoratorBehaviors ( type ) { let functionType ; if ( isFunction ( type ) ) { functionType = type ; } else if ( isOverloadedFunction ( type ) ) { functionType = type . overloads [ 0 ] ; } if ( ! functionType ) { return undefined ; } if ( functionType . details . decoratorDataClassBehaviors ) { return functionType . details . decoratorDataClassBehaviors ; } if ( functionType . details . fullName === 'dataclasses.dataclass' ) { return { keywordOnlyParams : false , generateEq : true , generateOrder : false , fieldDescriptorNames : [ 'dataclasses.field' , 'dataclasses.Field' ] , } ; } return undefined ; } function applyDataClassBehaviorOverride ( evaluator , errorNode , classType , argName , argValue ) { const fileInfo = AnalyzerNodeInfo . getFileInfo ( errorNode ) ; const value = evaluateStaticBoolExpression ( argValue , fileInfo . executionEnvironment ) ; switch ( argName ) { case 'order' : if ( value === true ) { classType . details . flags |= ClassTypeFlags . SynthesizedDataClassOrder ; } else if ( value === false ) { classType . details . flags &= ~ ClassTypeFlags . SynthesizedDataClassOrder ; } break ; case 'kw_only' : if ( value === false ) { classType . details . flags &= ~ ClassTypeFlags . DataClassKeywordOnlyParams ; } else if ( value === true ) { classType . details . flags |= ClassTypeFlags . DataClassKeywordOnlyParams ; } break ; case 'frozen' : { let hasUnfrozenBaseClass = false ; let hasFrozenBaseClass = false ; classType . details . baseClasses . forEach ( ( baseClass ) => { if ( isInstantiableClass ( baseClass ) && ClassType . isDataClass ( baseClass ) ) { if ( ClassType . isFrozenDataClass ( baseClass ) ) { hasFrozenBaseClass = true ; } else if ( ! baseClass . details . declaredMetaclass || ! isInstantiableClass ( baseClass . details . declaredMetaclass ) || ! baseClass . details . declaredMetaclass . details . metaclassDataClassTransform ) { hasUnfrozenBaseClass = true ; } } } ) ; if ( value === true || hasFrozenBaseClass ) { classType . details . flags |= ClassTypeFlags . FrozenDataClass ; if ( hasUnfrozenBaseClass ) { evaluator . addDiagnostic ( fileInfo . diagnosticRuleSet . reportGeneralTypeIssues , DiagnosticRule . reportGeneralTypeIssues , Localizer . Diagnostic . dataClassBaseClassNotFrozen ( ) , errorNode ) ; } } break ; } case 'init' : if ( value === false ) { classType . details . flags |= ClassTypeFlags . SkipSynthesizedDataClassInit ; } else if ( value === true ) { classType . details . flags &= ~ ClassTypeFlags . SkipSynthesizedDataClassInit ; } break ; case 'eq' : if ( value === false ) { classType . details . flags |= ClassTypeFlags . SkipSynthesizedDataClassEq ; } else if ( value === true ) { classType . details . flags &= ~ ClassTypeFlags . SkipSynthesizedDataClassEq ; } break ; case 'slots' : if ( value === true ) { classType . details . flags |= ClassTypeFlags . GenerateDataClassSlots ; if ( classType . details . localSlotsNames ) { evaluator . addDiagnostic ( fileInfo . diagnosticRuleSet . reportGeneralTypeIssues , DiagnosticRule . reportGeneralTypeIssues , Localizer . Diagnostic . dataClassSlotsOverwrite ( ) , errorNode ) ; } } else if ( value === false ) { classType . details . flags &= ~ ClassTypeFlags . GenerateDataClassSlots ; } break ; case 'hash' : case 'unsafe_hash' : if ( value === true ) { classType . details . flags |= ClassTypeFlags . SynthesizeDataClassUnsafeHash ; } break ; } } export function applyDataClassMetaclassBehaviorOverrides ( evaluator , classType , args ) { args . forEach ( ( arg ) => { if ( arg . valueExpression && arg . name ) { applyDataClassBehaviorOverride ( evaluator , arg . name , classType , arg . name . value , arg . valueExpression ) ; } } ) ; } export function applyDataClassDefaultBehaviors ( classType , defaultBehaviors ) { classType . details . dataClassBehaviors = defaultBehaviors ; classType . details . flags |= ClassTypeFlags . DataClass ; if ( defaultBehaviors . keywordOnlyParams ) { classType . details . flags |= ClassTypeFlags . DataClassKeywordOnlyParams ; } if ( ! defaultBehaviors . generateEq ) { classType . details . flags |= ClassTypeFlags . SkipSynthesizedDataClassEq ; } if ( defaultBehaviors . generateOrder ) { classType . details . flags |= ClassTypeFlags . SynthesizedDataClassOrder ; } } export function applyDataClassDecorator ( evaluator , classType , defaultBehaviors , callNode ) { applyDataClassDefaultBehaviors ( classType , defaultBehaviors ) ; if ( callNode ?. arguments ) { callNode . arguments . forEach ( ( arg ) => { if ( arg . name && arg . valueExpression ) { applyDataClassBehaviorOverride ( evaluator , arg , classType , arg . name . value , arg . valueExpression ) ; } } ) ; } }