1616using System . Collections . Generic ;
1717using System . Diagnostics ;
1818using System . Linq ;
19+ using System . Threading ;
1920using Microsoft . Python . Analysis . Analyzer . Evaluation ;
2021using Microsoft . Python . Analysis . Documents ;
2122using Microsoft . Python . Analysis . Modules ;
@@ -32,30 +33,36 @@ namespace Microsoft.Python.Analysis.Analyzer {
3233 internal class ModuleWalker : AnalysisWalker {
3334 private const string AllVariableName = "__all__" ;
3435 private readonly IDocumentAnalysis _stubAnalysis ;
36+ private readonly CancellationToken _cancellationToken ;
3537
3638 // A hack to use __all__ export in the most simple case.
3739 private int _allReferencesCount ;
3840 private bool _allIsUsable = true ;
3941
40- public ModuleWalker ( IServiceContainer services , IPythonModule module , PythonAst ast )
42+ public ModuleWalker ( IServiceContainer services , IPythonModule module , PythonAst ast , CancellationToken cancellationToken )
4143 : base ( new ExpressionEval ( services , module , ast ) ) {
4244 _stubAnalysis = Module . Stub is IDocument doc ? doc . GetAnyAnalysis ( ) : null ;
45+ _cancellationToken = cancellationToken ;
4346 }
4447
4548 public override bool Walk ( NameExpression node ) {
4649 if ( Eval . CurrentScope == Eval . GlobalScope && node . Name == AllVariableName ) {
4750 _allReferencesCount ++ ;
4851 }
52+
53+ _cancellationToken . ThrowIfCancellationRequested ( ) ;
4954 return base . Walk ( node ) ;
5055 }
5156
5257 public override bool Walk ( AugmentedAssignStatement node ) {
5358 HandleAugmentedAllAssign ( node ) ;
59+ _cancellationToken . ThrowIfCancellationRequested ( ) ;
5460 return base . Walk ( node ) ;
5561 }
5662
5763 public override bool Walk ( CallExpression node ) {
5864 HandleAllAppendExtend ( node ) ;
65+ _cancellationToken . ThrowIfCancellationRequested ( ) ;
5966 return base . Walk ( node ) ;
6067 }
6168
@@ -146,6 +153,7 @@ private bool IsHandleableAll(Node node) {
146153
147154 public override bool Walk ( PythonAst node ) {
148155 Check . InvalidOperation ( ( ) => Ast == node , "walking wrong AST" ) ;
156+ _cancellationToken . ThrowIfCancellationRequested ( ) ;
149157
150158 // Collect basic information about classes and functions in order
151159 // to correctly process forward references. Does not determine
@@ -181,16 +189,20 @@ public override bool Walk(PythonAst node) {
181189
182190 // Classes and functions are walked by their respective evaluators
183191 public override bool Walk ( ClassDefinition node ) {
192+ _cancellationToken . ThrowIfCancellationRequested ( ) ;
184193 SymbolTable . Evaluate ( node ) ;
185194 return false ;
186195 }
187196
188197 public override bool Walk ( FunctionDefinition node ) {
198+ _cancellationToken . ThrowIfCancellationRequested ( ) ;
189199 SymbolTable . Evaluate ( node ) ;
190200 return false ;
191201 }
192202
193203 public void Complete ( ) {
204+ _cancellationToken . ThrowIfCancellationRequested ( ) ;
205+
194206 SymbolTable . EvaluateAll ( ) ;
195207 SymbolTable . ReplacedByStubs . Clear ( ) ;
196208 MergeStub ( ) ;
@@ -220,11 +232,13 @@ public void Complete() {
220232 /// of the definitions. Stub may contains those so we need to merge it in.
221233 /// </remarks>
222234 private void MergeStub ( ) {
223- if ( Module . ModuleType == ModuleType . User ) {
235+ _cancellationToken . ThrowIfCancellationRequested ( ) ;
236+
237+ if ( Module . ModuleType == ModuleType . User || Module . ModuleType == ModuleType . Stub ) {
224238 return ;
225239 }
226240 // No stub, no merge.
227- if ( _stubAnalysis == null ) {
241+ if ( _stubAnalysis . IsEmpty ( ) ) {
228242 return ;
229243 }
230244 // TODO: figure out why module is getting analyzed before stub.
@@ -248,6 +262,18 @@ private void MergeStub() {
248262 if ( stubType . DeclaringModule is TypingModule && stubType . Name == "Any" ) {
249263 continue ;
250264 }
265+
266+ if ( sourceVar ? . Source == VariableSource . Import &&
267+ sourceVar . GetPythonType ( ) ? . DeclaringModule . Stub != null ) {
268+ // Keep imported types as they are defined in the library. For example,
269+ // 'requests' imports NullHandler as 'from logging import NullHandler'.
270+ // But 'requests' also declares NullHandler in its stub (but not in the main code)
271+ // and that declaration does not have documentation or location. Therefore avoid
272+ // taking types that are stub-only when similar type is imported from another
273+ // module that also has a stub.
274+ continue ;
275+ }
276+
251277 TryReplaceMember ( v , sourceType , stubType ) ;
252278 }
253279
@@ -257,7 +283,7 @@ private void MergeStub() {
257283 private void TryReplaceMember ( IVariable v , IPythonType sourceType , IPythonType stubType ) {
258284 // If type does not exist in module, but exists in stub, declare it unless it is an import.
259285 // If types are the classes, take class from the stub, then add missing members.
260- // Otherwise, replace type from one from the stub.
286+ // Otherwise, replace type by one from the stub.
261287 switch ( sourceType ) {
262288 case null :
263289 // Nothing in sources, but there is type in the stub. Declare it.
@@ -379,28 +405,26 @@ private static void TransferDocumentationAndLocation(IPythonType s, IPythonType
379405 if ( s != d && s is PythonType src && d is PythonType dst ) {
380406 // If type is a class, then doc can either come from class definition node of from __init__.
381407 // If class has doc from the class definition, don't stomp on it.
382- var transferDoc = true ;
383408 if ( src is PythonClassType srcClass && dst is PythonClassType dstClass ) {
384409 // Higher lever source wins
385410 if ( srcClass . DocumentationSource == PythonClassType . ClassDocumentationSource . Class ||
386411 ( srcClass . DocumentationSource == PythonClassType . ClassDocumentationSource . Init && dstClass . DocumentationSource == PythonClassType . ClassDocumentationSource . Base ) ) {
387412 dstClass . SetDocumentation ( srcClass . Documentation ) ;
388- transferDoc = false ;
389413 }
390- }
391-
392- // Sometimes destination (stub type) already has documentation. This happens when stub type
393- // is used to augment more than one type. For example, in threading module RLock stub class
394- // replaces both RLock function and _RLock class making 'factory' function RLock to look
395- // like a class constructor. Effectively a single stub type is used for more than one type
396- // in the source and two source types may have different documentation. Thus transferring doc
397- // from one source type affects documentation of another type. It may be better to clone stub
398- // type and separate instances for separate source type, but for now we'll just avoid stomping
399- // on the existing documentation.
400- if ( transferDoc && string . IsNullOrEmpty ( dst . Documentation ) ) {
401- var srcDocumentation = src . Documentation ;
402- if ( ! string . IsNullOrEmpty ( srcDocumentation ) ) {
403- dst . SetDocumentation ( srcDocumentation ) ;
414+ } else {
415+ // Sometimes destination (stub type) already has documentation. This happens when stub type
416+ // is used to augment more than one type. For example, in threading module RLock stub class
417+ // replaces both RLock function and _RLock class making 'factory' function RLock to look
418+ // like a class constructor. Effectively a single stub type is used for more than one type
419+ // in the source and two source types may have different documentation. Thus transferring doc
420+ // from one source type affects documentation of another type. It may be better to clone stub
421+ // type and separate instances for separate source type, but for now we'll just avoid stomping
422+ // on the existing documentation.
423+ if ( string . IsNullOrEmpty ( dst . Documentation ) ) {
424+ var srcDocumentation = src . Documentation ;
425+ if ( ! string . IsNullOrEmpty ( srcDocumentation ) ) {
426+ dst . SetDocumentation ( srcDocumentation ) ;
427+ }
404428 }
405429 }
406430
0 commit comments