Fixed #164: Type analyzer wasn't correctly inferring type of set when a set comprehension was used.

server/src/analyzer/expressionEvaluator.ts

@@ -2521,11 +2521,20 @@ export class ExpressionEvaluator {
     }
 
     private _getTypeFromSetExpression(node: SetNode): TypeResult {
-        const entryTypes = node.entries.map(expr => this._getTypeFromExpression(expr));
+        const entryTypes: Type[] = [];
 
         // Infer the set type based on the entries.
+        node.entries.forEach(entryNode => {
+            if (entryNode instanceof ListComprehensionNode) {
+                const setEntryType = this._getElementTypeFromListComprehensionExpression(entryNode);
+                entryTypes.push(setEntryType);
+            } else {
+                entryTypes.push(this._getTypeFromExpression(entryNode).type);
+            }
+        });
+
         const inferredEntryType = entryTypes.length > 0 ?
-            TypeUtils.combineTypes(entryTypes.map(e => e.type)) :
+            TypeUtils.combineTypes(entryTypes) :
             UnknownType.create();
 
         let type = ScopeUtils.getBuiltInObject(this._scope, 'set', [inferredEntryType]);

server/src/tests/samples/setComprehension1.py

@@ -0,0 +1,31 @@
+# This sample tests type checking for set comprehensions.
+
+from typing import Generator, Set
+
+a = [1, 2, 3, 4]
+
+def func1() -> Generator[int]:
+    b = (elem for elem in a)
+    return b
+
+
+def func2() -> Set[int]:
+    c = {elem for elem in a}
+    return c
+
+def func3() -> Set[str]:
+    c = {elem for elem in a}
+
+    # This should generate an error because
+    # c is a Set[int], which doesn't match
+    # the declared return type.
+    return c
+
+def generate():
+    for i in range(2):
+        yield i
+
+# Verify that generate returns a Generator.
+s = generate()
+s.close()
+

server/src/tests/typeAnalyzer.test.ts

@@ -397,6 +397,12 @@ test('ListComprehension1', () => {
     validateResults(analysisResults, 1);
 });
 
+test('SetComprehension1', () => {
+    let analysisResults = TestUtils.typeAnalyzeSampleFiles(['setComprehension1.py']);
+
+    validateResults(analysisResults, 1);
+});
+
 test('Literals1', () => {
     let analysisResults = TestUtils.typeAnalyzeSampleFiles(['literals1.py']);