Support $PSScriptRoot in dot-sourced files

Author: dee-see

src/PowerShellEditorServices/Language/AstOperations.cs

@@ -320,12 +320,12 @@ static private bool IsPowerShellDataFileAstNode(dynamic node, Type[] levelAstMap
         /// <summary>
         /// Finds all files dot sourced in a script
         /// </summary>
-        /// <param name="scriptAst">The abstract syntax tree of the given script</param>
+        /// <param name="scriptFile">The script file to use to find dot sourced files.</param>
         /// <returns></returns>
-        static public string[] FindDotSourcedIncludes(Ast scriptAst)
+        static public string[] FindDotSourcedIncludes(ScriptFile scriptFile)
         {
-            FindDotSourcedVisitor dotSourcedVisitor = new FindDotSourcedVisitor();
-            scriptAst.Visit(dotSourcedVisitor);
+            FindDotSourcedVisitor dotSourcedVisitor = new FindDotSourcedVisitor(scriptFile.FilePath);
+            scriptFile.ScriptAst.Visit(dotSourcedVisitor);
 
             return dotSourcedVisitor.DotSourcedFiles.ToArray();
         }

src/PowerShellEditorServices/Language/FindDotSourcedVisitor.cs

@@ -3,8 +3,11 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+using System;
 using System.Collections.Generic;
+using System.IO;
 using System.Management.Automation.Language;
+using System.Text.RegularExpressions;
 using Microsoft.PowerShell.EditorServices.Utility;
 
 namespace Microsoft.PowerShell.EditorServices
@@ -14,14 +17,17 @@ namespace Microsoft.PowerShell.EditorServices
     /// </summary>
     internal class FindDotSourcedVisitor : AstVisitor
     {
+        private readonly string _scriptDirectory;
+
         /// <summary>
         /// A hash set of the dot sourced files (because we don't want duplicates)
         /// </summary>
         public HashSet<string> DotSourcedFiles { get; private set; }
 
-        public FindDotSourcedVisitor()
+        public FindDotSourcedVisitor(string scriptPath)
         {
-            this.DotSourcedFiles = new HashSet<string>();
+            DotSourcedFiles = new HashSet<string>(StringComparer.CurrentCultureIgnoreCase);
+            _scriptDirectory = Path.GetDirectoryName(scriptPath);
         }
 
         /// <summary>
@@ -33,15 +39,44 @@ public FindDotSourcedVisitor()
         /// or a decision to continue if it wasn't found</returns>
         public override AstVisitAction VisitCommand(CommandAst commandAst)
         {
-            if (commandAst.InvocationOperator.Equals(TokenKind.Dot) &&
-                commandAst.CommandElements[0] is StringConstantExpressionAst)
+            CommandElementAst commandElementAst = commandAst.CommandElements[0];
+            if (commandAst.InvocationOperator.Equals(TokenKind.Dot))
             {
-                // Strip any quote characters off of the string
-                string fileName = PathUtils.NormalizePathSeparators(commandAst.CommandElements[0].Extent.Text.Trim('\'', '"'));
-                DotSourcedFiles.Add(fileName);
+                if (commandElementAst is StringConstantExpressionAst stringConstantExpressionAst)
+                {
+                    // Strip any quote characters off of the string
+                    DotSourcedFiles.Add(PathUtils.NormalizePathSeparators(stringConstantExpressionAst.Value));
+                }
+                else if (commandElementAst is ExpandableStringExpressionAst expandableStringExpressionAst)
+                {
+                    var path = GetPathFromExpandableStringExpression(expandableStringExpressionAst);
+                    if (path != null)
+                    {
+                        DotSourcedFiles.Add(PathUtils.NormalizePathSeparators(path));
+                    }
+                }
             }
 
             return base.VisitCommand(commandAst);
         }
+
+        private string GetPathFromExpandableStringExpression(ExpandableStringExpressionAst expandableStringExpressionAst)
+        {
+            var path = expandableStringExpressionAst.Value;
+            foreach (var nestedExpression in expandableStringExpressionAst.NestedExpressions)
+            {
+                if (nestedExpression is VariableExpressionAst variableExpressionAst
+                    && variableExpressionAst.VariablePath.UserPath.Equals("PSScriptRoot", StringComparison.CurrentCultureIgnoreCase))
+                {
+                    path = path.Replace(variableExpressionAst.ToString(), _scriptDirectory);
+                }
+                else
+                {
+                    return null; // We're going to get an invalid path anyway.
+                }
+            }
+
+            return path;
+        }
     }
 }

src/PowerShellEditorServices/Language/LanguageService.cs

@@ -14,6 +14,7 @@
 using System.Management.Automation.Language;
 using System.Runtime.InteropServices;
 using System.Security;
+using System.Text.RegularExpressions;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -426,11 +427,10 @@ public async Task<GetDefinitionResult> GetDefinitionOfSymbol(
                 if (foundSymbol.SymbolType == SymbolType.Function)
                 {
                     // Dot-sourcing is parsed as a "Function" Symbol.
-                    string trimmedName = PathUtils.NormalizePathSeparators(foundSymbol.SymbolName.Trim('\'', '"'));
-                    string dotSourcedPath = workspace.ResolveRelativeScriptPath(Path.GetDirectoryName(scriptFile.FilePath), trimmedName);
+                    string dotSourcedPath = GetDotSourcedPath(foundSymbol, workspace, scriptFile);
                     if (scriptFile.FilePath == dotSourcedPath)
                     {
-                        foundDefinition = new SymbolReference(SymbolType.Function, trimmedName, scriptFile.ScriptAst.Extent, scriptFile.FilePath);
+                        foundDefinition = new SymbolReference(SymbolType.Function, foundSymbol.SymbolName, scriptFile.ScriptAst.Extent, scriptFile.FilePath);
                         break;
                     }
                 }
@@ -487,6 +487,20 @@ await CommandHelpers.GetCommandInfo(
                 null;
         }
 
+        /// <summary>
+        /// Gets a path from a dot-source symbol.
+        /// </summary>
+        /// <param name="symbol">The symbol representing the dot-source expression.</param>
+        /// <param name="workspace">The current workspace</param>
+        /// <param name="scriptFile">The script file containing the symbol</param>
+        /// <returns></returns>
+        private static string GetDotSourcedPath(SymbolReference symbol, Workspace workspace, ScriptFile scriptFile)
+        {
+            string cleanedUpSymbol = PathUtils.NormalizePathSeparators(symbol.SymbolName.Trim('\'', '"'));
+            return workspace.ResolveRelativeScriptPath(Path.GetDirectoryName(scriptFile.FilePath),
+                Regex.Replace(cleanedUpSymbol, @"\$PSScriptRoot", Path.GetDirectoryName(scriptFile.FilePath), RegexOptions.IgnoreCase));
+        }
+
         /// <summary>
         /// Finds all the occurences of a symbol in the script given a file location
         /// </summary>

src/PowerShellEditorServices/Workspace/ScriptFile.cs

@@ -650,7 +650,7 @@ private void ParseFileContents()
 
             //Get all dot sourced referenced files and store  them
             this.ReferencedFiles =
-                AstOperations.FindDotSourcedIncludes(this.ScriptAst);
+                AstOperations.FindDotSourcedIncludes(this);
         }
 
         #endregion

test/PowerShellEditorServices.Test.Shared/References/DotSources.ps1

@@ -1,4 +1,5 @@
 . ./ReferenceFileE.ps1
+. "$PSScriptRoot/ReferenceFileE.ps1"
 . './ReferenceFileE.ps1'
 . "./ReferenceFileE.ps1"
 . .\ReferenceFileE.ps1

test/PowerShellEditorServices.Test.Shared/References/ReferenceFileB.ps1

@@ -1,5 +1,5 @@
-. .\ReferenceFileC.ps1
+. "$PSScriptRoot\ReferenceFileC.ps1"
 
 Get-ChildItem
 
-My-Function "testb"
+My-Function "testb"

test/PowerShellEditorServices.Test/Language/LanguageServiceTests.cs

@@ -172,7 +172,7 @@ await this.GetDefinition(
                 "Unexpected reference file: " + definitionResult.FoundDefinition.FilePath);
             Assert.Equal(1, definition.ScriptRegion.StartLineNumber);
             Assert.Equal(1, definition.ScriptRegion.StartColumnNumber);
-            Assert.Equal(PathUtils.NormalizePathSeparators("./ReferenceFileE.ps1"), definition.SymbolName);
+            Assert.Equal("./ReferenceFileE.ps1", definition.SymbolName);
         }
 
         [Fact]