microsoft · lostmsu · Nov 1, 2018 · Nov 1, 2018 · Nov 1, 2018 · Nov 1, 2018
diff --git a/src/.editorconfig b/src/.editorconfig
@@ -77,6 +77,8 @@ csharp_new_line_before_members_in_anonymous_types = false
 
 csharp_new_line_before_members_in_object_initializers = false
 
+csharp_prefer_braces = true:error
+
 dotnet_style_qualification_for_field = false:suggestion
 dotnet_style_qualification_for_property = false:none
 dotnet_style_qualification_for_method = false:none
diff --git a/src/Analysis/Engine/Impl/AnalysisUnit.cs b/src/Analysis/Engine/Impl/AnalysisUnit.cs
@@ -59,6 +59,9 @@ internal AnalysisUnit(ScopeStatement ast, IScope scope)
         internal AnalysisUnit(Node ast, PythonAst tree, IScope scope, bool forEval) {
             Ast = ast;
             Tree = tree;
+            if (ast == tree) {
+                _externalAnnotationAnalysisUnit = new Lazy<AnalysisUnit>(GetExternalAnalysisUnitImpl);
+            }
             _scope = scope;
             ForEval = forEval;
         }
@@ -342,6 +345,18 @@ public ILocationInfo ResolveLocation(object location) {
         internal virtual ILocationResolver AlternateResolver => null;
 
         ILocationResolver ILocationResolver.GetAlternateResolver() => AlternateResolver;
+
+        readonly Lazy<AnalysisUnit> _externalAnnotationAnalysisUnit;
+        protected internal virtual AnalysisUnit GetExternalAnnotationAnalysisUnit() => _externalAnnotationAnalysisUnit?.Value;
+
+        private AnalysisUnit GetExternalAnalysisUnitImpl() {
+            string analysisModuleName = ProjectEntry.ModuleName + PythonAnalyzer.AnnotationsModuleSuffix;
+            if (!ProjectEntry.ProjectState.Modules.TryImport(analysisModuleName, out var analysisModuleReference)) {
+                return null;
+            }
+
+            return analysisModuleReference.AnalysisModule?.AnalysisUnit;
+        }
     }
 
     class ClassAnalysisUnit : AnalysisUnit {
@@ -462,6 +477,23 @@ internal static IAnalysisSet EvaluateBaseClass(DDG ddg, ClassInfo newClass, int
 
             return bases;
         }
+
+        protected internal override AnalysisUnit GetExternalAnnotationAnalysisUnit() {
+            var parentAnnotation = _outerUnit.GetExternalAnnotationAnalysisUnit();
+            if (parentAnnotation == null) {
+                return null;
+            }
+
+            if (!parentAnnotation.Scope.TryGetVariable(Ast.Name, out var annotationVariable)) {
+                return null;
+            }
+
+            if (annotationVariable.Types is ClassInfo classAnnotation) {
+                return classAnnotation.AnalysisUnit;
+            }
+
+            return (annotationVariable.Types.OnlyOneOrDefault() as ClassInfo)?.AnalysisUnit;
+        }
     }
 
     class ComprehensionAnalysisUnit : AnalysisUnit {

diff --git a/src/Analysis/Engine/Impl/Analyzer/DDG.cs b/src/Analysis/Engine/Impl/Analyzer/DDG.cs
@@ -113,12 +113,25 @@ public override bool Walk(PythonAst node) {
             if (!ProjectState.Modules.TryImport(_unit.DeclaringModule.Name, out existingRef)) {
                 // publish our module ref now so that we don't collect dependencies as we'll be fully processed
                 if (existingRef == null) {
-                    ProjectState.Modules[_unit.DeclaringModule.Name] = new ModuleReference(_unit.DeclaringModule, _unit.DeclaringModule.Name);
+                    ProjectState.Modules[_unit.DeclaringModule.Name] = existingRef = new ModuleReference(_unit.DeclaringModule, _unit.DeclaringModule.Name);
                 } else {
                     existingRef.Module = _unit.DeclaringModule;
                 }
             }
 
+            string annotationModuleName = _unit.DeclaringModule.Name + PythonAnalyzer.AnnotationsModuleSuffix;
+            if (ProjectState.Modules.TryGetImportedModule(annotationModuleName, out var _)
+                && TryImportModule(annotationModuleName, true, out var annotationsReference, out var _)) {
+                annotationsReference.Module.Imported(_unit);
+
+                if (ProjectState.Modules.TryGetImportedModule(annotationModuleName, out annotationsReference)) {
+                    FinishImportModuleOrMember(annotationsReference, attribute: null, name: "__pyi__",
+                        addRef: true, node: node, nameReference: new NameExpression("__pyi__"));
+                } else {
+                    Debug.Fail($"Failed to get module {annotationModuleName} we just imported");
+                }
+            }
+
             return base.Walk(node);
         }
 
@@ -345,19 +358,23 @@ private bool AssignImportedModuleOrMember(string name, IAnalysisSet value, bool
         /// True to add <paramref name="node"/> as a reference of the
         /// imported value.
         /// </param>
-        private void FinishImportModuleOrMember(ModuleReference module, IReadOnlyList<string> attribute, string name, bool addRef, Node node, NameExpression nameReference) {
-            if (AssignImportedModuleOrMember(
+        private bool FinishImportModuleOrMember(ModuleReference module, IReadOnlyList<string> attribute, string name, bool addRef, Node node, NameExpression nameReference) {
+            bool imported = AssignImportedModuleOrMember(
                 name,
                 GetImportedModuleOrMember(module, attribute, addRef, node, nameReference, attribute?.Count == 1 ? name : null),
                 addRef,
                 node,
                 nameReference
-            )) {
+            );
+
+            if (imported) {
                 // Imports into our global scope need to enqueue modules that have imported us
                 if (Scope == GlobalScope.Scope) {
                     GlobalScope.ModuleDefinition.EnqueueDependents();
                 }
             }
+
+            return imported;
         }
 
         public override bool Walk(FromImportStatement node) {

diff --git a/src/Analysis/Engine/Impl/Analyzer/FunctionAnalysisUnit.cs b/src/Analysis/Engine/Impl/Analyzer/FunctionAnalysisUnit.cs
@@ -72,6 +72,23 @@ internal override ModuleInfo GetDeclaringModule() {
             return base.GetDeclaringModule() ?? _declUnit.DeclaringModule;
         }
 
+        protected internal override AnalysisUnit GetExternalAnnotationAnalysisUnit() {
+            var parentAnnotation = _declUnit.GetExternalAnnotationAnalysisUnit();
+            if (parentAnnotation == null) {
+                return null;
+            }
+
+            if (!parentAnnotation.Scope.TryGetVariable(Ast.Name, out var annotationVariable)) {
+                return null;
+            }
+
+            if (annotationVariable.Types is FunctionInfo functionAnnotation) {
+                return functionAnnotation.AnalysisUnit;
+            }
+
+            return (annotationVariable.Types.OnlyOneOrDefault() as FunctionInfo)?.AnalysisUnit;
+        }
+
         internal override void AnalyzeWorker(DDG ddg, CancellationToken cancel) {
             // Resolve default parameters and decorators in the outer scope but
             // continue to associate changes with this unit.
@@ -184,16 +201,31 @@ internal IAnalysisSet ProcessFunctionDecorators(DDG ddg) {
         internal void AnalyzeDefaultParameters(DDG ddg) {
             IVariableDefinition param;
             var scope = (FunctionScope)Scope;
+            var annotationAnalysis = GetExternalAnnotationAnalysisUnit() as FunctionAnalysisUnit;
+            var functionAnnotation = annotationAnalysis?.Function.FunctionDefinition;
+            if (functionAnnotation?.Parameters.Length != Ast.Parameters.Length) {
+                functionAnnotation = null;
+            }
             for (var i = 0; i < Ast.Parameters.Length; ++i) {
                 var p = Ast.Parameters[i];
-                if (p.Annotation != null) {
-                    var val = ddg._eval.EvaluateAnnotation(p.Annotation);
-                    if (val?.Any() == true && Scope.TryGetVariable(p.Name, out param)) {
-                        param.AddTypes(this, val, false);
-                        var vd = scope.GetParameter(p.Name);
-                        if (vd != null && vd != param) {
-                            vd.AddTypes(this, val, false);
-                        }
+                var annotation = p.Annotation;
+                IAnalysisSet annotationValue = null;
+                if (annotation != null) {
+                    annotationValue = ddg._eval.EvaluateAnnotation(annotation);
+                } else if (functionAnnotation?.Parameters[i].Annotation != null) {
+                    try {
+                        ddg.SetCurrentUnit(annotationAnalysis);
+                        annotationValue = ddg._eval.EvaluateAnnotation(functionAnnotation.Parameters[i].Annotation);
+                    } finally {
+                        ddg.SetCurrentUnit(this);
+                    }
+                }
+
+                if (annotationValue?.Any() == true && Scope.TryGetVariable(p.Name, out param)) {
+                    param.AddTypes(this, annotationValue, false);
+                    var vd = scope.GetParameter(p.Name);
+                    if (vd != null && vd != param) {
+                        vd.AddTypes(this, annotationValue, false);
                     }
                 }
 
@@ -209,8 +241,20 @@ internal void AnalyzeDefaultParameters(DDG ddg) {
                     }
                 }
             }
+
+            IAnalysisSet ann = null;
             if (Ast.ReturnAnnotation != null) {
-                var ann = ddg._eval.EvaluateAnnotation(Ast.ReturnAnnotation);
+                ann = ddg._eval.EvaluateAnnotation(Ast.ReturnAnnotation);
+            } else if (functionAnnotation?.ReturnAnnotation != null) {
+                try {
+                    ddg.SetCurrentUnit(annotationAnalysis);
+                    ann = ddg._eval.EvaluateAnnotation(functionAnnotation.ReturnAnnotation);
+                } finally {
+                    ddg.SetCurrentUnit(this);
+                }
+            }
+
+            if (ann != null) {
                 var resType = ann;
                 if (Ast.IsGenerator) {
                     if (ann.Split<ProtocolInfo>(out var gens, out resType)) {

diff --git a/src/Analysis/Engine/Impl/Infrastructure/Extensions/EnumerableExtensions.cs b/src/Analysis/Engine/Impl/Infrastructure/Extensions/EnumerableExtensions.cs
@@ -60,6 +60,24 @@ public static bool SetEquals<T>(this IEnumerable<T> source, IEnumerable<T> other
         public static IEnumerable<T> Keys<T, U>(this IEnumerable<KeyValuePair<T, U>> source) => source.Select(GetKey);
         public static IEnumerable<T> ExcludeDefault<T>(this IEnumerable<T> source) => source.Where(i => !Equals(i, default(T)));
 
+        /// <summary>
+        /// Returns the only element of sequence, or the default value of <typeparamref name="T"/>, if sequence has &lt;&gt; 1 elements.
+        /// </summary>
+        public static T OnlyOneOrDefault<T>(this IEnumerable<T> source) {
+            using (var enumerator = source.GetEnumerator()) {
+                if (!enumerator.MoveNext()) {
+                    return default(T);
+                }
+
+                T result = enumerator.Current;
+
+                if (enumerator.MoveNext()) {
+                    return default(T);
+                }
+
+                return result;
+            }
+        }
 
         public static IEnumerable<T> TraverseBreadthFirst<T>(this T root, Func<T, IEnumerable<T>> selectChildren) {
             var items = new Queue<T>();

diff --git a/src/Analysis/Engine/Impl/ProjectEntry.cs b/src/Analysis/Engine/Impl/ProjectEntry.cs
@@ -46,7 +46,7 @@ internal sealed class ProjectEntry : IPythonProjectEntry, IAggregateableProjectE
         private readonly HashSet<AggregateProjectEntry> _aggregates = new HashSet<AggregateProjectEntry>();
 
         private TaskCompletionSource<IModuleAnalysis> _analysisTcs = new TaskCompletionSource<IModuleAnalysis>();
-        private AnalysisUnit _unit;
+        internal AnalysisUnit _unit;
         private readonly ManualResetEventSlim _pendingParse = new ManualResetEventSlim(true);
         private long _expectedParseVersion;
         private long _expectedAnalysisVersion;

diff --git a/src/Analysis/Engine/Impl/PythonAnalyzer.cs b/src/Analysis/Engine/Impl/PythonAnalyzer.cs
@@ -36,6 +36,7 @@ namespace Microsoft.PythonTools.Analysis {
     /// </summary>
     public partial class PythonAnalyzer : IPythonAnalyzer, IDisposable {
         public const string PythonAnalysisSource = "Python";
+        internal const string AnnotationsModuleSuffix = "__pyi__";
         private static object _nullKey = new object();
 
         private readonly bool _disposeInterpreter;
@@ -216,6 +217,21 @@ public IPythonProjectEntry AddModule(string moduleName, string filePath, Uri doc
             return entry;
         }
 
+        /// <summary>
+        /// Adds a new annotations file to the list of available annotation files and returns a ProjectEntry object.
+        /// 
+        /// This method is thread safe.
+        /// </summary>
+        /// <param name="moduleName">The name of the module; used to associate with imports</param>
+        /// <param name="filePath">The path to the file on disk</param>
+        /// <param name="cookie">An application-specific identifier for the module</param>
+        /// <returns>The project entry for the new module.</returns>
+        public IPythonProjectEntry AddModuleAnnotations(string moduleName, string filePath, Uri documentUri = null, IAnalysisCookie cookie = null) {
+            Check.ArgumentNotNull(nameof(moduleName), moduleName);
+
+            return AddModule(moduleName + AnnotationsModuleSuffix, filePath, documentUri, cookie);
+        }
+
         /// <summary>
         /// Associates an existing module with a new name.
         /// </summary>

diff --git a/src/Analysis/Engine/Impl/Values/ModuleInfo.cs b/src/Analysis/Engine/Impl/Values/ModuleInfo.cs
@@ -123,6 +123,8 @@ public override IDictionary<string, IAnalysisSet> GetAllMembers(IModuleContext m
 
         public ModuleInfo ParentPackage { get; set; }
 
+        public override AnalysisUnit AnalysisUnit => _projectEntry._unit;
+
         public void AddChildPackage(ModuleInfo childPackage, AnalysisUnit curUnit, string realName = null) {
             realName = realName ?? childPackage.Name;
             int lastDot;