Added support for @type_check_only decorator. This addresses #5810.

packages/pyright-internal/src/analyzer/decorators.ts

@@ -89,6 +89,8 @@ export function getFunctionFlagsFromDecorators(evaluator: TypeEvaluator, node: F
                 flags |= FunctionTypeFlags.Final;
             } else if (decoratorType.details.builtInName === 'override') {
                 flags |= FunctionTypeFlags.Overridden;
+            } else if (decoratorType.details.builtInName === 'type_check_only') {
+                flags |= FunctionTypeFlags.TypeCheckOnly;
             }
         } else if (isInstantiableClass(decoratorType)) {
             if (ClassType.isBuiltIn(decoratorType, 'staticmethod')) {
@@ -187,6 +189,11 @@ export function applyFunctionDecorator(
             return inputFunctionType;
         }
 
+        if (decoratorType.details.builtInName === 'type_check_only') {
+            undecoratedType.details.flags |= FunctionTypeFlags.TypeCheckOnly;
+            return inputFunctionType;
+        }
+
         // Handle property setters and deleters.
         if (decoratorNode.expression.nodeType === ParseNodeType.MemberAccess) {
             const baseType = evaluator.getTypeOfExpression(
@@ -353,6 +360,11 @@ export function applyClassDecorator(
             return inputClassType;
         }
 
+        if (decoratorType.details.builtInName === 'type_check_only') {
+            originalClassType.details.flags |= ClassTypeFlags.TypeCheckOnly;
+            return inputClassType;
+        }
+
         if (decoratorType.details.builtInName === 'deprecated') {
             originalClassType.details.deprecatedMessage = getCustomDeprecationMessage(decoratorNode);
             return inputClassType;
@@ -401,13 +413,6 @@ function getTypeOfDecorator(evaluator: TypeEvaluator, node: DecoratorNode, funct
 
     const decoratorTypeResult = evaluator.getTypeOfExpression(node.expression, flags);
 
-    // Special-case typing.type_check_only. It's used in many stdlib
-    // functions, and pyright treats it as a no-op, so don't waste
-    // time evaluating it.
-    if (isFunction(decoratorTypeResult.type) && decoratorTypeResult.type.details.builtInName === 'type_check_only') {
-        return functionOrClassType;
-    }
-
     // Special-case the combination of a classmethod decorator applied
     // to a property. This is allowed in Python 3.9, but it's not reflected
     // in the builtins.pyi stub for classmethod.

packages/pyright-internal/src/analyzer/typeEvaluator.ts

@@ -1228,6 +1228,34 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
         return typeResult;
     }
 
+    // Reports the case where a function or class has been decorated with
+    // @type_check_only and is used in a value expression.
+    function reportUseOfTypeCheckOnly(type: Type, node: ExpressionNode) {
+        let isTypeCheckingOnly = false;
+        let name = '';
+
+        if (isInstantiableClass(type) && !type.includeSubclasses) {
+            isTypeCheckingOnly = ClassType.isTypeCheckOnly(type);
+            name = type.details.name;
+        } else if (isFunction(type)) {
+            isTypeCheckingOnly = FunctionType.isTypeCheckOnly(type);
+            name = type.details.name;
+        }
+
+        if (isTypeCheckingOnly) {
+            const fileInfo = AnalyzerNodeInfo.getFileInfo(node);
+
+            if (!fileInfo.isStubFile) {
+                addDiagnostic(
+                    AnalyzerNodeInfo.getFileInfo(node).diagnosticRuleSet.reportGeneralTypeIssues,
+                    DiagnosticRule.reportGeneralTypeIssues,
+                    Localizer.Diagnostic.typeCheckOnly().format({ name }),
+                    node
+                );
+            }
+        }
+    }
+
     function reportInvalidUseOfPep695TypeAlias(type: Type, node: ExpressionNode): boolean {
         // PEP 695 type aliases cannot be used as instantiable classes.
         if (type.typeAliasInfo?.name && type.typeAliasInfo.isPep695Syntax && TypeBase.isSpecialForm(type)) {
@@ -4422,6 +4450,10 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
             }
         }
 
+        if ((flags & EvaluatorFlags.ExpectingTypeAnnotation) === 0) {
+            reportUseOfTypeCheckOnly(type, node);
+        }
+
         if ((flags & EvaluatorFlags.DisallowPep695TypeAlias) !== 0) {
             if (reportInvalidUseOfPep695TypeAlias(type, node)) {
                 type = UnknownType.create();
@@ -5087,7 +5119,7 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
                     return { type: UnknownType.create(isIncomplete), isIncomplete };
                 }
 
-                if (flags & EvaluatorFlags.ExpectingTypeAnnotation) {
+                if ((flags & EvaluatorFlags.ExpectingTypeAnnotation) !== 0) {
                     if (!isIncomplete) {
                         addDiagnostic(
                             AnalyzerNodeInfo.getFileInfo(node).diagnosticRuleSet.reportGeneralTypeIssues,
@@ -5401,6 +5433,10 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
             type = isFunctionRule ? AnyType.create() : UnknownType.create();
         }
 
+        if ((flags & EvaluatorFlags.ExpectingTypeAnnotation) === 0) {
+            reportUseOfTypeCheckOnly(type, node.memberName);
+        }
+
         // Should we specialize the class?
         if ((flags & EvaluatorFlags.DoNotSpecialize) === 0) {
             if (isInstantiableClass(type) && !type.typeArguments) {
@@ -6764,7 +6800,7 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
                             MemberAccessFlags.SkipAttributeAccessOverride | MemberAccessFlags.ConsiderMetaclassOnly
                         );
 
-                        if (flags & EvaluatorFlags.ExpectingTypeAnnotation) {
+                        if ((flags & EvaluatorFlags.ExpectingTypeAnnotation) !== 0) {
                             // If the class doesn't derive from Generic, a type argument should not be allowed.
                             addDiagnostic(
                                 AnalyzerNodeInfo.getFileInfo(node).diagnosticRuleSet.reportGeneralTypeIssues,
@@ -15107,7 +15143,11 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
                 // methods that are abstract are overridden and shouldn't
                 // cause the TypedDict to be marked as abstract.
                 if (isInstantiableClass(baseClass) && ClassType.isBuiltIn(baseClass, '_TypedDict')) {
-                    baseClass.details.flags &= ~ClassTypeFlags.SupportsAbstractMethods;
+                    baseClass = ClassType.cloneWithNewFlags(
+                        baseClass,
+                        baseClass.details.flags &
+                            ~(ClassTypeFlags.SupportsAbstractMethods | ClassTypeFlags.TypeCheckOnly)
+                    );
                 }
             }
         }

packages/pyright-internal/src/analyzer/types.ts

@@ -548,6 +548,9 @@ export const enum ClassTypeFlags {
 
     // For dataclasses, should __hash__ be generated?
     SynthesizeDataClassUnsafeHash = 1 << 26,
+
+    // Decorated with @type_check_only.
+    TypeCheckOnly = 1 << 27,
 }
 
 export interface DataClassBehaviors {
@@ -845,6 +848,13 @@ export namespace ClassType {
         return newClassType;
     }
 
+    export function cloneWithNewFlags(classType: ClassType, newFlags: ClassTypeFlags): ClassType {
+        const newClassType = TypeBase.cloneType(classType);
+        newClassType.details = { ...newClassType.details };
+        newClassType.details.flags = newFlags;
+        return newClassType;
+    }
+
     export function isLiteralValueSame(type1: ClassType, type2: ClassType): boolean {
         if (type1.literalValue === undefined) {
             return type2.literalValue === undefined;
@@ -994,6 +1004,10 @@ export namespace ClassType {
         return !!(classType.details.flags & ClassTypeFlags.SynthesizeDataClassUnsafeHash);
     }
 
+    export function isTypeCheckOnly(classType: ClassType) {
+        return !!(classType.details.flags & ClassTypeFlags.TypeCheckOnly);
+    }
+
     export function isTypedDictClass(classType: ClassType) {
         return !!(classType.details.flags & ClassTypeFlags.TypedDictClass);
     }
@@ -1281,6 +1295,9 @@ export const enum FunctionTypeFlags {
     // for implied methods such as those used in namedtuple, dataclass, etc.
     SynthesizedMethod = 1 << 6,
 
+    // Decorated with @type_check_only.
+    TypeCheckOnly = 1 << 7,
+
     // Function is decorated with @overload
     Overloaded = 1 << 8,
 
@@ -1890,6 +1907,10 @@ export namespace FunctionType {
         return (type.details.flags & FunctionTypeFlags.SynthesizedMethod) !== 0;
     }
 
+    export function isTypeCheckOnly(type: FunctionType): boolean {
+        return (type.details.flags & FunctionTypeFlags.TypeCheckOnly) !== 0;
+    }
+
     export function isOverloaded(type: FunctionType): boolean {
         return (type.details.flags & FunctionTypeFlags.Overloaded) !== 0;
     }

packages/pyright-internal/src/localization/localize.ts

@@ -896,6 +896,8 @@ export namespace Localizer {
                 getRawString('Diagnostic.typeAssignmentMismatchWildcard')
             );
         export const typeCallNotAllowed = () => getRawString('Diagnostic.typeCallNotAllowed');
+        export const typeCheckOnly = () =>
+            new ParameterizedString<{ name: string }>(getRawString('Diagnostic.typeCheckOnly'));
         export const typeCommentDeprecated = () => getRawString('Diagnostic.typeCommentDeprecated');
         export const typedDictAccess = () => getRawString('Diagnostic.typedDictAccess');
         export const typedDictBadVar = () => getRawString('Diagnostic.typedDictBadVar');

packages/pyright-internal/src/localization/package.nls.en-us.json

@@ -441,6 +441,7 @@
         "typeAssignmentMismatch": "Expression of type \"{sourceType}\" cannot be assigned to declared type \"{destType}\"",
         "typeAssignmentMismatchWildcard": "Import symbol \"{name}\" has type \"{sourceType}\", which cannot be assigned to declared type \"{destType}\"",
         "typeCallNotAllowed": "type() call should not be used in type annotation",
+        "typeCheckOnly": "\"{name}\" is marked as @type_check_only and can be used only in type annotations",
         "typeCommentDeprecated": "Use of type comments is deprecated; use type annotation instead",
         "typedDictAccess": "Could not access item in TypedDict",
         "typedDictBadVar": "TypedDict classes can contain only type annotations",

packages/pyright-internal/src/tests/samples/typeCheckOnly1.py

@@ -0,0 +1,40 @@
+# This sample tests the reporting of a function or class decorated with
+# @type_check_only when used in a value expression.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, type_check_only
+
+if TYPE_CHECKING:
+    from typing import _TypedDict
+
+a1: function
+a2: _TypedDict
+
+# This should generate an error.
+v1 = function
+
+# This should generate an error.
+v2 = isinstance(1, ellipsis)
+
+# This should generate an error.
+v3 = _TypedDict
+
+
+if TYPE_CHECKING:
+
+    class ClassA:
+        @type_check_only
+        def method1(self):
+            pass
+
+    @type_check_only
+    def func1() -> None:
+        ...
+
+
+# This should generate an error.
+ClassA().method1()
+
+# This should generate an error.
+func1()

packages/pyright-internal/src/tests/samples/typeNarrowingIsEllipsis1.py

@@ -6,7 +6,7 @@
 from typing import TypeVar
 
 
-_T = TypeVar("_T", str, ellipsis)
+_T = TypeVar("_T", str, type(Ellipsis))
 
 
 def func1(val: int | ellipsis):

packages/pyright-internal/src/tests/typeEvaluator3.test.ts

@@ -1691,33 +1691,3 @@ test('Decorator7', () => {
 
     TestUtils.validateResults(analysisResults, 0);
 });
-
-test('DataclassTransform1', () => {
-    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['dataclassTransform1.py']);
-
-    TestUtils.validateResults(analysisResults, 6);
-});
-
-test('DataclassTransform2', () => {
-    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['dataclassTransform2.py']);
-
-    TestUtils.validateResults(analysisResults, 6);
-});
-
-test('DataclassTransform3', () => {
-    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['dataclassTransform3.py']);
-
-    TestUtils.validateResults(analysisResults, 6);
-});
-
-test('DataclassTransform4', () => {
-    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['dataclassTransform4.py']);
-
-    TestUtils.validateResults(analysisResults, 1);
-});
-
-test('Async1', () => {
-    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['async1.py']);
-
-    TestUtils.validateResults(analysisResults, 6);
-});

packages/pyright-internal/src/tests/typeEvaluator5.test.ts

@@ -317,3 +317,38 @@ test('TypedDictReadOnly2', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typedDictReadOnly2.py'], configOptions);
     TestUtils.validateResults(analysisResults, 8);
 });
+
+test('DataclassTransform1', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['dataclassTransform1.py']);
+
+    TestUtils.validateResults(analysisResults, 6);
+});
+
+test('DataclassTransform2', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['dataclassTransform2.py']);
+
+    TestUtils.validateResults(analysisResults, 6);
+});
+
+test('DataclassTransform3', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['dataclassTransform3.py']);
+
+    TestUtils.validateResults(analysisResults, 6);
+});
+
+test('DataclassTransform4', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['dataclassTransform4.py']);
+
+    TestUtils.validateResults(analysisResults, 1);
+});
+
+test('Async1', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['async1.py']);
+
+    TestUtils.validateResults(analysisResults, 6);
+});
+
+test('TypeCheckOnly1', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeCheckOnly1.py']);
+    TestUtils.validateResults(analysisResults, 5);
+});