From 6fbd8ba1dd8777760567179cda875f1c2605d108 Mon Sep 17 00:00:00 2001
From: MartinGC94 <42123497+MartinGC94@users.noreply.github.com>
Date: Fri, 28 Feb 2025 05:20:25 +0100
Subject: [PATCH] Handle type inference for redirected commands (#21131)

---
 .../engine/parser/TypeInferenceVisitor.cs     | 67 +++++++++++++++++++
 .../engine/Api/TypeInference.Tests.ps1        | 41 ++++++++++++
 2 files changed, 108 insertions(+)

diff --git a/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs b/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
index 0e0355761a..d010a48853 100644
--- a/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
+++ b/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
@@ -1272,6 +1272,73 @@ namespace System.Management.Automation
 
         private void InferTypesFrom(CommandAst commandAst, List<PSTypeName> inferredTypes)
         {
+            if (commandAst.Redirections.Count > 0)
+            {
+                var mergedStreams = new HashSet<RedirectionStream>();
+                bool allStreamsMerged = false;
+                foreach (RedirectionAst streamRedirection in commandAst.Redirections)
+                {
+                    if (streamRedirection is FileRedirectionAst fileRedirection)
+                    {
+                        if (fileRedirection.FromStream is RedirectionStream.All or RedirectionStream.Output)
+                        {
+                            // command output is redirected so it returns nothing.
+                            return;
+                        }
+                    }
+                    else if (streamRedirection is MergingRedirectionAst mergeRedirection && mergeRedirection.ToStream == RedirectionStream.Output)
+                    {
+                        if (mergeRedirection.FromStream == RedirectionStream.All)
+                        {
+                            allStreamsMerged = true;
+                            continue;
+                        }
+
+                        _ = mergedStreams.Add(mergeRedirection.FromStream);
+                    }
+                }
+
+                if (allStreamsMerged)
+                {
+                    inferredTypes.Add(new PSTypeName(typeof(ErrorRecord)));
+                    inferredTypes.Add(new PSTypeName(typeof(WarningRecord)));
+                    inferredTypes.Add(new PSTypeName(typeof(VerboseRecord)));
+                    inferredTypes.Add(new PSTypeName(typeof(DebugRecord)));
+                    inferredTypes.Add(new PSTypeName(typeof(InformationRecord)));
+                }
+                else
+                {
+                    foreach (RedirectionStream value in mergedStreams)
+                    {
+                        switch (value)
+                        {
+                            case RedirectionStream.Error:
+                                inferredTypes.Add(new PSTypeName(typeof(ErrorRecord)));
+                                break;
+
+                            case RedirectionStream.Warning:
+                                inferredTypes.Add(new PSTypeName(typeof(WarningRecord)));
+                                break;
+
+                            case RedirectionStream.Verbose:
+                                inferredTypes.Add(new PSTypeName(typeof(VerboseRecord)));
+                                break;
+
+                            case RedirectionStream.Debug:
+                                inferredTypes.Add(new PSTypeName(typeof(DebugRecord)));
+                                break;
+
+                            case RedirectionStream.Information:
+                                inferredTypes.Add(new PSTypeName(typeof(InformationRecord)));
+                                break;
+
+                            default:
+                                break;
+                        }
+                    }
+                }
+            }
+
             PseudoBindingInfo pseudoBinding = new PseudoParameterBinder()
             .DoPseudoParameterBinding(commandAst, null, null, PseudoParameterBinder.BindingType.ParameterCompletion);
 
diff --git a/test/powershell/engine/Api/TypeInference.Tests.ps1 b/test/powershell/engine/Api/TypeInference.Tests.ps1
index 83bab06d19..39bdeedbad 100644
--- a/test/powershell/engine/Api/TypeInference.Tests.ps1
+++ b/test/powershell/engine/Api/TypeInference.Tests.ps1
@@ -1492,6 +1492,47 @@ Describe "Type inference Tests" -tags "CI" {
         $null = [AstTypeInference]::InferTypeOf($FoundAst)
     }
 
+    It 'Infers type of command with all streams redirected to Success stream' {
+        $res = [AstTypeInference]::InferTypeOf( { Get-PSDrive *>&1 }.Ast)
+        $ExpectedTypeNames = @(
+            [ErrorRecord].FullName
+            [WarningRecord].FullName
+            [VerboseRecord].FullName
+            [DebugRecord].FullName
+            [InformationRecord].FullName
+            [PSDriveInfo].FullName
+        ) -join ';'
+        $res.Name -join ';' | Should -Be $ExpectedTypeNames
+    }
+
+    It 'Infers type of command with success stream redirected' {
+        $res = [AstTypeInference]::InferTypeOf( { Get-PSDrive *>&1 1>$null }.Ast)
+        $res.Count | Should -Be 0
+    }
+
+    It 'Infers type of command with some streams redirected to success' {
+        $res = [AstTypeInference]::InferTypeOf( { Get-PSDrive 3>&1 4>&1 }.Ast)
+        $res.Count | Should -Be 3
+        $ExpectedTypeNames = @(
+            [PSDriveInfo].FullName
+            [VerboseRecord].FullName
+            [WarningRecord].FullName
+        ) -join ';'
+        ($res.Name | Sort-Object) -join ';' | Should -Be $ExpectedTypeNames
+    }
+
+    It 'Infers type of command with other streams redirected to success' {
+        $res = [AstTypeInference]::InferTypeOf( { Get-PSDrive 2>&1 5>&1 6>&1 }.Ast)
+        $res.Count | Should -Be 4
+        $ExpectedTypeNames = @(
+            [DebugRecord].FullName
+            [ErrorRecord].FullName
+            [InformationRecord].FullName
+            [PSDriveInfo].FullName
+        ) -join ';'
+        ($res.Name | Sort-Object) -join ';' | Should -Be $ExpectedTypeNames
+    }
+
     It 'Should only consider assignments wrapped in parentheses to be a part of the output in a Named block' {
         $res = [AstTypeInference]::InferTypeOf( { [string]$Assignment1 = "Hello"; ([int]$Assignment2 = 42) }.Ast)
         $res.Count | Should -Be 1