From 0a5f6bc3532d6d541a8bf694a3a4a17384747430 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Fri, 27 Oct 2023 15:15:41 +0100
Subject: [PATCH] [daemon] Fix return type change to optional in generic
 function

Previously changing a return type to an optional type was not propagated
at least in some cases, since astdiff could simplify away the optional
type.
---
 mypy/server/astdiff.py   |  4 +++-
 test-data/unit/diff.test | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/mypy/server/astdiff.py b/mypy/server/astdiff.py
index 93f178dca35a2..5323bf2c57cbe 100644
--- a/mypy/server/astdiff.py
+++ b/mypy/server/astdiff.py
@@ -74,6 +74,7 @@ class level -- these are handled at attribute level (say, 'mod.Cls.method'
     Var,
 )
 from mypy.semanal_shared import find_dataclass_transform_spec
+from mypy.state import state
 from mypy.types import (
     AnyType,
     CallableType,
@@ -456,7 +457,8 @@ def normalize_callable_variables(self, typ: CallableType) -> CallableType:
                 tv = v.copy_modified(id=tid)
             tvs.append(tv)
             tvmap[v.id] = tv
-        return expand_type(typ, tvmap).copy_modified(variables=tvs)
+        with state.strict_optional_set(True):
+            return expand_type(typ, tvmap).copy_modified(variables=tvs)
 
     def visit_tuple_type(self, typ: TupleType) -> SnapshotItem:
         return ("TupleType", snapshot_types(typ.items))
diff --git a/test-data/unit/diff.test b/test-data/unit/diff.test
index 66adfaecd909d..8fc74868123e8 100644
--- a/test-data/unit/diff.test
+++ b/test-data/unit/diff.test
@@ -1497,3 +1497,36 @@ class C:
     def meth(self) -> int: return 0
 [out]
 __main__.C.meth
+
+[case testGenericFunctionWithOptionalReturnType]
+from typing import Type, TypeVar
+
+T = TypeVar("T")
+
+class C:
+    @classmethod
+    def get_by_team_and_id(
+        cls: Type[T],
+        raw_member_id: int,
+        include_removed: bool = False,
+    ) -> T:
+        pass
+
+[file next.py]
+from typing import Type, TypeVar, Optional
+
+T = TypeVar("T")
+
+class C:
+    @classmethod
+    def get_by_team_and_id(
+        cls: Type[T],
+        raw_member_id: int,
+        include_removed: bool = False,
+    ) -> Optional[T]:
+        pass
+
+[builtins fixtures/classmethod.pyi]
+[out]
+__main__.C.get_by_team_and_id
+__main__.Optional