Merge pull request proxb#157 from copdips/dev

fix proxb#156 - Start-RSJob - be able to import private functions
MVKozlov · Jan 27, 2018 · a0ee818 · a0ee818
2 parents 07c9778 + 6131382
commit a0ee818
Show file tree

Hide file tree

Showing 7 changed files with 150 additions and 32 deletions.
diff --git a/PoshRSJob/PoshRSJob.psd1 b/PoshRSJob/PoshRSJob.psd1
@@ -84,8 +84,8 @@ AliasesToExport = 'gsj','rmsj','rsj','spsj','ssj','wsj'
 #ModuleList = @()
 
 # List of all files packaged with this module
-FileList = 'PoshRSJob.psd1', 'PoshRSJob.psm1', 'en-US\about_PoshRSJob.help.txt', 'Private\ConvertScript.ps1', 'Private\ConvertScriptBlockV2.ps1', 
-'Private\FindFunction.ps1', 'Private\GetParamVariable.ps1', 'Private\GetUsingVariables.ps1', 'Private\GetUsingVariablesV2.ps1',
+FileList = 'PoshRSJob.psd1', 'PoshRSJob.psm1', 'en-US\about_PoshRSJob.help.txt', 'Private\ConvertScript.ps1', 'Private\ConvertScriptBlockV2.ps1',
+'Private\FindFunction.ps1', 'Private\GetFunctionByFile.ps1', 'Private\GetFunctionDefinitionByFunction.ps1', 'Private\GetParamVariable.ps1', 'Private\GetUsingVariables.ps1', 'Private\GetUsingVariablesV2.ps1',
 'Private\Increment.ps1', 'Private\RegisterScriptScopeFunction.ps1', 'Public\Get-RSJob.ps1', 'Public\Receive-RSJob.ps1',
 'Public\Remove-RSJob.ps1', 'Public\Start-RSJob.ps1', 'Public\Stop-RSJob.ps1', 'Public\Wait-RSJob.ps1', 'TypeData\PoshRSJob.Format.ps1xml', 'TypeData\PoshRSJob.Types.ps1xml',
 'Private\SetIsReceived.ps1'
@@ -95,28 +95,27 @@ PrivateData = @{
     PSData = @{
 			# The primary categorization of this module (from the TechNet Gallery tech tree).
 			Category = "Multithreading"
-			
+
 			# Keyword tags to help users find this module via navigations and search.
 			Tags = @('PoshRSJob', 'Runspace','RunspacePool', 'Linux', 'PowerShellCore', 'RSJob')
-			
+
 			# The web address of an icon which can be used in galleries to represent this module
 			#IconUri = ''
-			
+
 			# The web address of this module's project or support homepage.
 			ProjectUri = "https://github.com/proxb/PoshRSJob"
-			
+
 			# The web address of this module's license. Points to a page that's embeddable and linkable.
 			LicenseUri = "https://opensource.org/licenses/MIT"
-			
+
 			# Release notes for this particular version of the module
 			# ReleaseNotes = False
-			
+
 			# If true, the LicenseUrl points to an end-user license (not just a source license) which requires the user agreement before use.
 			RequireLicenseAcceptance = "False"
-			
+
 			# Indicates this is a pre-release/testing version of the module.
 			IsPrerelease = 'False'
 		}
     }
 }
-
diff --git a/PoshRSJob/Private/FindFunction.ps1 b/PoshRSJob/Private/FindFunction.ps1
@@ -3,7 +3,7 @@ Function FindFunction {
     [CmdletBinding()]
     param (
         [string]$ScriptBlock
-    )     
+    )
     #Just in case we have some oddness going on
     $ScriptBlock = $ScriptBlock -replace '`','``'
     # Tokenize the script
@@ -25,22 +25,22 @@ Function FindFunction {
         }
         if ($functionsearch) {
             If ($i -gt 1 -AND ($tokens[$i].StartLine -eq $tokens[$i-1].EndLine)) {
-                $SpaceCount = $tokens[$i].StartColumn - $tokens[$i-1].EndColumn            
+                $SpaceCount = $tokens[$i].StartColumn - $tokens[$i-1].EndColumn
                 $space = ' '*"$($SpaceCount)"
                 If ($SpaceCount -gt 0) {
                     If ($SpaceCount -notmatch '^[5|9]$') {
-                        Write-Verbose "Adding Space: $($SpaceCount)" 
+                        Write-Verbose "Adding Space: $($SpaceCount)"
                         [void]$Definition.Append($Space)
                     } ElseIf ($SpaceCount -match '^[5|9]$') {
-                        Write-Verbose "Adding NewLine" 
+                        Write-Verbose "Adding NewLine"
                         [void]$Definition.Append("`n")
                     }
                 }
             }
-            Write-Verbose $tokens[$i].Content 
+            Write-Verbose $tokens[$i].Content
             Switch ($tokens[$i].Type) {
                 'NewLine' {
-                    Write-Verbose 'Adding NewLine' 
+                    Write-Verbose 'Adding NewLine'
                     [void]$Definition.Append("`n")
                 }
                 'CommandArgument' {
@@ -71,17 +71,17 @@ Function FindFunction {
                     [void]$Definition.Append($tokens[$i].Content)
                 }
                 'Variable' {
-                    [void]$Definition.Append("`$$($tokens[$i].Content)")                      
+                    [void]$Definition.Append("`$$($tokens[$i].Content)")
                 }
                 'Type' {
                     Switch ($PSVersionTable.PSVersion.Major) {
                         '2' {
-                            [void]$Definition.Append("[$($tokens[$i].Content)]")                         
+                            [void]$Definition.Append("[$($tokens[$i].Content)]")
                         }
                         Default {
-                            [void]$Definition.Append($($tokens[$i].Content)) 
+                            [void]$Definition.Append($($tokens[$i].Content))
                         }
-                    }                                         
+                    }
                 }
                 Default {
                     [void]$Definition.Append($tokens[$i].Content)
@@ -97,4 +97,4 @@ Function FindFunction {
             }
         }
     }
-}
+}
diff --git a/PoshRSJob/Private/GetFunctionByFile.ps1 b/PoshRSJob/Private/GetFunctionByFile.ps1
@@ -0,0 +1,35 @@
+Function GetFunctionByFile {
+    [CmdletBinding()]
+    param (
+        [string[]]$FilePath
+    )
+
+    $psMajorVersion = $PSVersionTable.PSVersion.Major
+    $functionsInFile = @()
+    ForEach ($thisFilePath in $FilePath) {
+        Write-Verbose "Working on file : $thisFilePath"
+
+        if (-not (Test-Path $thisFilePath)) {
+            Write-Warning "Cannot find file : $thisFilePath"
+            continue
+        }
+
+        try {
+            Switch ($psMajorVersion) {
+                '2' {
+                    $scriptBlockInFile = [ScriptBlock]::Create($(Get-Content $thisFilePath) -join [Environment]::NewLine)
+                    $functionsInFile += @(FindFunction -ScriptBlock $scriptBlockInFile)
+                }
+                Default {
+                    $AST = [System.Management.Automation.Language.Parser]::ParseFile($thisFilePath, [ref]$null, [ref]$null)
+                    $functionsInFile += $AST.FindAll( {$args[0] -is [System.Management.Automation.Language.FunctionDefinitionAst]} , $true)
+                }
+            }
+            Write-Verbose "Functions found in file : $($functionsInFile.Name -join '; ')"
+        }
+        catch {
+            Write-Warning "$thisFilePath : $($_.Exception.Message)"
+        }
+    }
+    $functionsInFile
+}
diff --git a/PoshRSJob/Private/GetFunctionDefinitionByFunction.ps1 b/PoshRSJob/Private/GetFunctionDefinitionByFunction.ps1
@@ -0,0 +1,16 @@
+Function GetFunctionDefinitionByFunction {
+    [CmdletBinding()]
+    param (
+        [parameter(ValueFromPipeline = $True)]
+        $FunctionItem
+    )
+
+    if ($FunctionItem -is [PSCustomObject]) {
+        # In case of Powershell v2
+        $function.Body.Trim().Trim("{}")
+    }
+    else {
+        # In case of Powershell v3+
+        $function.Body.Extent.Text.Trim("{}")
+    }
+}
diff --git a/PoshRSJob/Private/GetUsingVariables.ps1 b/PoshRSJob/Private/GetUsingVariables.ps1
@@ -1,4 +1,4 @@
 Function GetUsingVariables {
     Param ([scriptblock]$ScriptBlock)
-    $ScriptBlock.ast.FindAll({$args[0] -is [System.Management.Automation.Language.UsingExpressionAst]},$True)    
+    $ScriptBlock.ast.FindAll( {$args[0] -is [System.Management.Automation.Language.UsingExpressionAst]}, $True)
 }
diff --git a/PoshRSJob/Public/Start-RSJob.ps1 b/PoshRSJob/Public/Start-RSJob.ps1
@@ -42,6 +42,9 @@ Function Start-RSJob {
         .PARAMETER FunctionsToImport
             A collection of functions that will be imported for use with a background runspace job.
 
+        .PARAMETER FunctionFilesToImport
+            A collection of files containing custom functions that will be imported into the background runspace job.
+
         .PARAMETER VariablesToImport
             A collection of variables that will be imported for use with a background runspace job.
             If used, $using:variable not expanded !
@@ -160,11 +163,11 @@ Function Start-RSJob {
         DefaultParameterSetName = 'ScriptBlock'
     )]
     Param (
-        [parameter(Mandatory=$True,Position=0,ParameterSetName = 'ScriptBlock')]
+        [parameter(Mandatory = $True, Position = 0, ParameterSetName = 'ScriptBlock')]
         [ScriptBlock]$ScriptBlock,
-        [parameter(Position=0,ParameterSetName = 'ScriptPath')]
+        [parameter(Position = 0, ParameterSetName = 'ScriptPath')]
         [string]$FilePath,
-        [parameter(ValueFromPipeline=$True,ValueFromPipelineByPropertyName=$True)]
+        [parameter(ValueFromPipeline = $True, ValueFromPipelineByPropertyName = $True)]
         [object]$InputObject,
         [parameter()]
         [object]$Name,
@@ -184,41 +187,52 @@ Function Start-RSJob {
         [Alias('FunctionsToLoad')]
         [string[]]$FunctionsToImport,
         [parameter()]
+        [Alias('FunctionFilesToLoad')]
+        [string[]]$FunctionFilesToImport,
+        [parameter()]
         [Alias('VariablesToLoad')]
         [string[]]$VariablesToImport
     )
     Begin {
+
         If ($PSBoundParameters['Debug']) {
             $DebugPreference = 'Continue'
         }
+
         Write-Debug "[BEGIN]"
+
         If ($PSBoundParameters.ContainsKey('Verbose')) {
             Write-Verbose "Displaying PSBoundParameters"
             $PSBoundParameters.GetEnumerator() | ForEach-Object {
                 Write-Verbose $_
             }
         }
+
         If ($PSBoundParameters.ContainsKey('Name')) {
             If ($Name -isnot [scriptblock]) {
                 $JobName = [scriptblock]::Create("Write-Output `"$Name`"")
             }
             Else {
-                $JobName = [scriptblock]::Create( ($Name -replace '\$_','$Item'))
+                $JobName = [scriptblock]::Create( ($Name -replace '\$_', '$Item'))
             }
         }
         Else {
             Write-Verbose "Creating default Job Name"
             $JobName = [scriptblock]::Create('Write-Output Job$($Id)')
         }
+
         $InitialSessionState = [System.Management.Automation.Runspaces.InitialSessionState]::CreateDefault()
+
         If ($PSBoundParameters['ModulesToImport']) {
             [void]$InitialSessionState.ImportPSModule($ModulesToImport)
         }
+
         If ($PSBoundParameters['PSSnapinsToImport']) {
             ForEach ($PSSnapin in $PSSnapinsToImport) {
-                [void]$InitialSessionState.ImportPSSnapIn($PSSnapin,[ref]$Null)
+                [void]$InitialSessionState.ImportPSSnapIn($PSSnapin, [ref]$Null)
             }
         }
+
         If ($PSBoundParameters['FunctionsToImport']) {
             Write-Verbose "Loading custom functions: $($FunctionsToImport -join '; ')"
             ForEach ($Function in $FunctionsToImport) {
@@ -235,11 +249,42 @@ Function Start-RSJob {
 
                 #Check for an alias and add it as well
                 If ($Alias = Get-Alias | Where-Object { $_.Definition -eq $Function }) {
-                    $AliasEntry = New-Object System.Management.Automation.Runspaces.SessionStateAliasEntry -ArgumentList $Alias.Name,$Alias.Definition
+                    $AliasEntry = New-Object System.Management.Automation.Runspaces.SessionStateAliasEntry -ArgumentList $Alias.Name, $Alias.Definition
                     $InitialSessionState.Commands.Add($AliasEntry)
                 }
             }
         }
+
+        If ($PSBoundParameters['FunctionFilesToImport']) {
+            Write-Verbose "Loading custom function files : $($FunctionFilesToImport -join '; ')"
+            $functionsInFiles = GetFunctionByFile -FilePath $FunctionFilesToImport
+
+            if ($null -eq $functionsInFiles) {
+                Write-Warning "Cannot find any functions in given files"
+            }
+            else {
+                ForEach ($function in $functionsInFiles) {
+                    $functionName = $function.Name
+                    Write-Verbose "Loading custom function : $functionName"
+
+                    try {
+                        $functionDefinition = GetFunctionDefinitionByFunction -FunctionItem $function
+                        $SessionStateFunction = New-Object System.Management.Automation.Runspaces.SessionStateFunctionEntry -ArgumentList $functionName, $functionDefinition
+                        $InitialSessionState.Commands.Add($SessionStateFunction)
+                    }
+                    catch {
+                        Write-Warning "$($functionName): $($_.Exception.Message)"
+                    }
+
+                    #Check for an alias and add it as well
+                    if ($Alias = Get-Alias | Where-Object { $_.Definition -eq $Function }) {
+                        $AliasEntry = New-Object System.Management.Automation.Runspaces.SessionStateAliasEntry -ArgumentList $Alias.Name, $Alias.Definition
+                        $InitialSessionState.Commands.Add($AliasEntry)
+                    }
+                }
+            }
+        }
+
         If ($PSBoundParameters['VariablesToImport']) {
             Write-Verbose "Loading variables: $($VariablesToImport -join '; ')"
             $UserVariables = New-Object System.Collections.ArrayList
@@ -304,6 +349,7 @@ Function Start-RSJob {
         }
         Write-Debug "ListCount: $($List.Count)"
     }
+
     Process {
         Write-Debug "[PROCESS]"
         If ($PSBoundParameters.ContainsKey('InputObject')) {
@@ -312,6 +358,7 @@ Function Start-RSJob {
             $ForeachDetected = $false
         }
     }
+
     End {
         Write-Debug "[END]"
         $SBParamVars = @(GetParamVariable -ScriptBlock $ScriptBlock)
@@ -511,7 +558,7 @@ Function Start-RSJob {
             If ($UsingVariableValues.count -gt 0) {
                 For ($i=0;$i -lt $UsingVariableValues.count;$i++) {
                     Write-Verbose "Adding Param: $($UsingVariableValues[$i].Name) Value: $($UsingVariableValues[$i].Value)"
-                    [void]$PowerShell.AddParameter($UsingVariableValues[$i].NewVarName,$UsingVariableValues[$i].Value)
+                    [void]$PowerShell.AddParameter($UsingVariableValues[$i].NewVarName, $UsingVariableValues[$i].Value)
                 }
             }
             Write-Verbose "Checking for ArgumentList"
@@ -546,7 +593,7 @@ Function Start-RSJob {
                 Finished = $Handle.IsCompleted
                 Command  = $ScriptBlock.ToString()
                 RunspacePoolID = $RunSpacePoolID
-                Batch = $Batch
+                Batch          = $Batch
             }
 
             $RSPObject.LastActivity = Get-Date

diff --git a/Tests/PoshRSJob.Tests.ps1 b/Tests/PoshRSJob.Tests.ps1
@@ -207,16 +207,36 @@ Describe "Start-RSJob PS$PSVersion" {
             } ) | Wait-RSJob | Receive-RSJob
             $Output1 | Should Be 1
         }
+        It 'should support FunctionFilesToImport syntax' {
+            $functionFile1 = Join-Path $env:TEMP ([GUID]::NewGuid()).Guid
+            $functionFile2 = Join-Path $env:TEMP ([GUID]::NewGuid()).Guid
+            "# test multi-lines scriptblock with comment
+            function f1 {
+                Write-Output 'r1'
+            }
+            function f2 {Write-Output 'r2'}" | Out-File $functionFile1
+            "function f3 {Write-Output 'r3'} function f4 {Write-Output 'r4'}" | Out-File $functionFile2
+            $Output = @(
+                Start-RSJob  @Verbose -ScriptBlock {
+                    f1
+                    f2
+                    f3
+                    f4
+                } -FunctionFilesToImport $functionFile1, $functionFile2 | Wait-RSJob | Receive-RSJob)
+            ($Output -join ',') | Should Be 'r1,r2,r3,r4'
+            Remove-Item $functionFile1, $functionFile2
+        }
         It 'should support VariablesToImport syntax' {
             $Output2 = @(
                 $tester0 = 'tester012'; $testvar1 = 'testvar124'; $testvar2 = 'testvar248'
                 Start-RSJob  @Verbose -ScriptBlock {
                     $tester0
                     $testvar1
                     $testvar2
-                } -VariablesToImport tester0,testvar* | Wait-RSJob | Receive-RSJob)
+                } -VariablesToImport tester0, testvar* | Wait-RSJob | Receive-RSJob)
             ($Output2 -join ',') | Should Be 'tester012,testvar124,testvar248'
         }
+
     }
 }
 
@@ -227,7 +247,8 @@ Describe "Get-RSJob PS$PSVersion" {
             $Output = @( Get-RSJob @Verbose )
             $Props = $Output[0].PSObject.Properties | Select-Object -ExpandProperty Name
 
-            $Output.count | Should be 11
+            # Write-Output $Output
+            $Output.count | Should be 12
             $Props -contains "Id" | Should be $True
             $Props -contains "State" | Should be $True
             $Props -contains "HasMoreData" | Should be $True