More tests, handling nested generics

CTrando · CTrando · commit cb5ff36db2ff · 2019-08-15T11:14:29.000-07:00
diff --git a/src/Analysis/Ast/Impl/Analyzer/Evaluation/ExpressionEval.Generics.cs b/src/Analysis/Ast/Impl/Analyzer/Evaluation/ExpressionEval.Generics.cs
@@ -138,16 +138,19 @@ private IReadOnlyList<IMember> EvaluateIndex(IndexExpression expr) {
 
         /// <summary>
         /// Given an index argument, will try and resolve it to a forward reference, e.g
+        /// Forward references are types declared in quotes, e.g 'int' is equivalent to type int
         /// 
         /// List['str'] => List[str]
+        /// 'A[int]' => A[int]
         /// </summary>
         private IMember GetValueFromForwardRef(IMember index) {
-            index.TryGetConstant(out string memberName);
-            if (string.IsNullOrEmpty(memberName)) {
+            index.TryGetConstant(out string forwardRefStr);
+            if (string.IsNullOrEmpty(forwardRefStr)) {
                 return null;
             }
 
-            return LookupNameInScopes(memberName, out var _, out var _, LookupOptions.Normal);
+            var forwardRefExpr = TryCreateExpression(forwardRefStr);
+            return GetValueFromExpression(forwardRefExpr);
         }
 
         private IReadOnlyList<IMember> EvaluateCallArgs(CallExpression expr) {
diff --git a/src/Analysis/Ast/Test/GenericsTests.cs b/src/Analysis/Ast/Test/GenericsTests.cs
@@ -1079,32 +1079,34 @@ def get(self) -> _T:
         public async Task GenericForwardRef() {
             const string code = @"
 from typing import List, Dict
+
+x = List['A']
+y = Dict['A', 'B']
+
 class A:
     def test(self) -> int:
         pass
 class B:
     def test(self) -> int:
         pass
-l = List['A']
-d = Dict['A', 'B']
+z = List['A']
+w = Dict['A', 'B']
 ";
             var analysis = await GetAnalysisAsync(code, PythonVersions.LatestAvailable3X);
-            analysis.Should().HaveVariable("l")
+            analysis.Should().HaveVariable("x")
                 .Which.Should().HaveType("List[A]");
-            analysis.Should().HaveVariable("d")
+            analysis.Should().HaveVariable("y")
+                           .Which.Should().HaveType("Dict[A, B]");
+            analysis.Should().HaveVariable("z")
+                .Which.Should().HaveType("List[A]");
+            analysis.Should().HaveVariable("w")
                            .Which.Should().HaveType("Dict[A, B]");
         }
 
         [TestMethod, Priority(0)]
         public async Task GenericBuiltinTypeForwardRef() {
             const string code = @"
 from typing import List, Dict
-class A:
-    def test(self) -> int:
-        pass
-class B:
-    def test(self) -> int:
-        pass
 l = List['int']
 d = Dict['float', 'str']
 ";
@@ -1119,15 +1121,16 @@ def test(self) -> int:
         public async Task GenericListForwardRef() {
             const string code = @"
 from typing import List, Dict
-class A: ...
-class B: ...
 
 def test() -> 'List[A]':
     pass
 
 def test1() -> 'Dict[A, B]':
     pass
 
+class A: ...
+class B: ...
+
 x = test()
 y = test1()
 ";
@@ -1147,23 +1150,55 @@ public async Task GenericClassForwardRef() {
 T = TypeVar('T')
 K = TypeVar('K')
 
-class A(Generic[T]):
-    def test(self) -> T:
-        pass
-
 class B(Generic[T, K]):
     def test(self) -> T:
         pass
 
     def test1(self) -> K:
         pass
 
+b = B['A', 'A']()
+y = b.test()
+z = b.test1()
+
+class A(Generic[T]):
+    def test(self) -> T:
+        pass
+";
+            var analysis = await GetAnalysisAsync(code, PythonVersions.LatestAvailable3X);
+            analysis.Should().HaveVariable("b")
+                          .Which.Should().HaveType("B[A, A]");
+            analysis.Should().HaveVariable("y")
+                           .Which.Should().HaveType("A");
+            analysis.Should().HaveVariable("z")
+                                      .Which.Should().HaveType("A");
+        }
+
+        [TestMethod, Priority(0)]
+        public async Task GenericClassForwardRefNestedGenerics() {
+            const string code = @"
+from typing import Generic, TypeVar
+
+T = TypeVar('T')
+K = TypeVar('K')
+
 a = A['B']()
 x = a.test()
 
-b = B['A', 'A']()
+b = B['A[int]', 'A[str]']()
 y = b.test()
 z = b.test1()
+
+class A(Generic[T]):
+    def test(self) -> T:
+        pass
+
+class B(Generic[T, K]):
+    def test(self) -> T:
+        pass
+
+    def test1(self) -> K:
+        pass
 ";
             var analysis = await GetAnalysisAsync(code, PythonVersions.LatestAvailable3X);
             analysis.Should().HaveVariable("a")
@@ -1172,13 +1207,14 @@ def test1(self) -> K:
                            .Which.Should().HaveType("B");
 
             analysis.Should().HaveVariable("b")
-                          .Which.Should().HaveType("B[A, A]");
+                          .Which.Should().HaveType("B[A[int], A[str]]");
             analysis.Should().HaveVariable("y")
-                           .Which.Should().HaveType("A");
+                           .Which.Should().HaveType("A[int]");
             analysis.Should().HaveVariable("z")
-                                      .Which.Should().HaveType("A");
+                                      .Which.Should().HaveType("A[str]");
         }
 
+
         [TestMethod, Priority(0)]
         public async Task GenericFunctionArguments() {
             const string code = @"
