Add support for excluding functions from coverage with attribute (#2593)

* Add support for excluding functions with [ExcludeFromCodeCoverageAttribute()] attribute

* Refactor coverage exclusion logic

* Improving Test-ContainsAttribute and minor review improvements

* Replacing FindAll with AstVisitor

* Updating tests

* Adding description to new class and high level test

src/csharp/Pester/CoverageLocationVisitor.cs

@@ -0,0 +1,80 @@
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Management.Automation.Language;
+
+namespace Pester
+{
+    /// <summary>
+    /// A visitor class for traversing the PowerShell AST to collect coverage-relevant locations.
+    /// This replaces predicate-based filtering with a centralized, extensible approach.
+    ///
+    /// Advantages:
+    /// - Efficiently skips nodes with attributes like [ExcludeFromCodeCoverage].
+    /// - Simplifies logic by handling each AST type in dedicated methods.
+    /// </summary>
+    public class CoverageLocationVisitor : AstVisitor2
+    {
+        public readonly List<Ast> CoverageLocations = new();
+
+        public override AstVisitAction VisitScriptBlock(ScriptBlockAst scriptBlockAst)
+        {
+            if (scriptBlockAst.ParamBlock?.Attributes != null)
+            {
+                foreach (var attribute in scriptBlockAst.ParamBlock.Attributes)
+                {
+                    if (attribute.TypeName.GetReflectionType() == typeof(ExcludeFromCodeCoverageAttribute))
+                    {
+                        return AstVisitAction.SkipChildren;
+                    }
+                }
+            }
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitCommand(CommandAst commandAst)
+        {
+            CoverageLocations.Add(commandAst);
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitCommandExpression(CommandExpressionAst commandExpressionAst)
+        {
+            CoverageLocations.Add(commandExpressionAst);
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitDynamicKeywordStatement(DynamicKeywordStatementAst dynamicKeywordStatementAst)
+        {
+            CoverageLocations.Add(dynamicKeywordStatementAst);
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitBreakStatement(BreakStatementAst breakStatementAst)
+        {
+            CoverageLocations.Add(breakStatementAst);
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitContinueStatement(ContinueStatementAst continueStatementAst)
+        {
+            CoverageLocations.Add(continueStatementAst);
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitExitStatement(ExitStatementAst exitStatementAst)
+        {
+            CoverageLocations.Add(exitStatementAst);
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitThrowStatement(ThrowStatementAst throwStatementAst)
+        {
+            CoverageLocations.Add(throwStatementAst);
+            return AstVisitAction.Continue;
+        }
+
+        // ReturnStatementAst is excluded as it's not behaving consistent.
+        // "return" is not hit in 5.1 but fixed in a later version. Using "return 123" we get hit on 123 but not return.
+        // See https://github.com/pester/Pester/issues/1465#issuecomment-604323645
+    }
+}

src/functions/Coverage.ps1

@@ -284,28 +284,9 @@ function Get-CommandsInFile {
     $tokens = $null
     $ast = [System.Management.Automation.Language.Parser]::ParseFile($Path, [ref] $tokens, [ref] $errors)
 
-    if ($PSVersionTable.PSVersion.Major -ge 5) {
-        # In PowerShell 5.0, dynamic keywords for DSC configurations are represented by the DynamicKeywordStatementAst
-        # class.  They still trigger breakpoints, but are not a child class of CommandBaseAst anymore.
-
-        # ReturnStatementAst is excluded as it's not behaving consistent.
-        # "return" is not hit in 5.1 but fixed in a later version. Using "return 123" we get hit on 123 but not return.
-        # See https://github.com/pester/Pester/issues/1465#issuecomment-604323645
-        $predicate = {
-            $args[0] -is [System.Management.Automation.Language.DynamicKeywordStatementAst] -or
-            $args[0] -is [System.Management.Automation.Language.CommandBaseAst] -or
-            $args[0] -is [System.Management.Automation.Language.BreakStatementAst] -or
-            $args[0] -is [System.Management.Automation.Language.ContinueStatementAst] -or
-            $args[0] -is [System.Management.Automation.Language.ExitStatementAst] -or
-            $args[0] -is [System.Management.Automation.Language.ThrowStatementAst]
-        }
-    }
-    else {
-        $predicate = { $args[0] -is [System.Management.Automation.Language.CommandBaseAst] }
-    }
-
-    $searchNestedScriptBlocks = $true
-    $ast.FindAll($predicate, $searchNestedScriptBlocks)
+    $visitor = [Pester.CoverageLocationVisitor]::new()
+    $ast.Visit($visitor)
+    return $visitor.CoverageLocations
 }
 
 function Test-CoverageOverlapsCommand {

tst/functions/Coverage.Tests.ps1

@@ -1284,6 +1284,131 @@ InPesterModuleScope {
         }
     }
 
+    Describe 'Coverage Location Visitor' {
+        BeforeAll {
+            $testScript = @'
+            using namespace System.Diagnostics.CodeAnalysis
+
+            function FunctionIncluded {
+                "I am included"
+            }
+
+            function FunctionExcluded {
+                [ExcludeFromCodeCoverageAttribute()]
+                param()
+
+                "I am not included"
+            }
+
+            FunctionIncluded
+            FunctionExcluded
+
+'@
+            $tokens = $null
+            $errors = $null
+            $ast = [System.Management.Automation.Language.Parser]::ParseInput($testScript, [ref]$tokens, [ref]$errors)
+        }
+
+        Context 'Collect coverage locations' {
+            BeforeAll {
+                $visitor = [Pester.CoverageLocationVisitor]::new()
+                $ast.Visit($visitor)
+            }
+
+            It 'Skips excluded script blocks' {
+                $excludedCommand = $visitor.CoverageLocations | Where-Object {
+                    $_ -is [System.Management.Automation.Language.CommandExpressionAst] -and
+                    $_.Expression.Value -eq "I am not included"
+                }
+
+                $excludedCommand.Count | Should -Be 0 -Because "Command in excluded script blocks should not be collected."
+            }
+
+            It 'Processes included script blocks' {
+                $includedCommand = $visitor.CoverageLocations | Where-Object {
+                    $_ -is [System.Management.Automation.Language.CommandExpressionAst] -and
+                    $_.Expression.Value -eq "I am included"
+                }
+
+                $includedCommand.Count | Should -Be 1 -Because "Command in included script blocks should be collected."
+            }
+        }
+
+        Context 'Collect coverage locations for other AST types' {
+            It 'Collects all relevant AST types' {
+                $script = @'
+                foreach ($i in 1..10) { # 1 location
+                    break               # 1 location
+                    continue            # 1 location
+                    if ($i -eq 5) {     # 1 location
+                        throw           # 1 location
+                    }
+                    if ($i -eq 7) {     # 1 location
+                        exit            # 1 location
+                    }
+                    return              # not collected
+                }
+'@
+                $tokens = $null
+                $errors = $null
+                $ast = [System.Management.Automation.Language.Parser]::ParseInput($script, [ref]$tokens, [ref]$errors)
+
+                $visitor = [Pester.CoverageLocationVisitor]::new()
+                $ast.Visit($visitor)
+
+                $visitor.CoverageLocations.Count | Should -Be 7 -Because "Break, Continue, Throw, and Exit statements should be collected."
+            }
+        }
+
+        Context 'Coverage analysis with exclusion using <description>' -Foreach @(
+            @{ UseBreakpoints = $true; Description = "With breakpoints" }
+            @{ UseBreakpoints = $false; Description = "Profiler-based coverage collection" }
+        ) {
+            BeforeAll {
+                $root = (Get-PSDrive TestDrive).Root
+                $testScriptPath = Join-Path -Path $root -ChildPath TestScript.ps1
+                Set-Content -Path $testScriptPath -Value $testScript
+
+                $breakpoints = Enter-CoverageAnalysis -CodeCoverage @{ Path = $testScriptPath } -UseBreakpoints $UseBreakpoints
+
+                @($breakpoints).Count | Should -Be 3 -Because 'The correct number of breakpoints should be defined.'
+
+                if ($UseBreakpoints) {
+                    & $testScriptPath
+                }
+                else {
+                    $patched, $tracer = Start-TraceScript $breakpoints
+                    try { & $testScriptPath } finally { Stop-TraceScript -Patched $patched }
+                    $measure = $tracer.Hits
+                }
+
+                $coverageReport = Get-CoverageReport -CommandCoverage $breakpoints -Measure $measure
+            }
+
+            It 'Correctly reports executed commands' {
+                $coverageReport.NumberOfCommandsExecuted | Should -Be 3 -Because 'The executed commands count should match.'
+            }
+
+            It 'Correctly reports analyzed commands' {
+                $coverageReport.NumberOfCommandsAnalyzed | Should -Be 3 -Because 'All commands should be analyzed.'
+            }
+
+            It 'Correctly reports missed commands' {
+                $coverageReport.MissedCommands.Count | Should -Be 0 -Because 'No command should be missed.'
+            }
+
+            It 'Correctly reports hit commands' {
+                $coverageReport.HitCommands.Count | Should -Be 3 -Because 'Three commands should be hit.'
+            }
+
+            AfterAll {
+                if ($UseBreakpoints) {
+                    Exit-CoverageAnalysis -CommandCoverage $breakpoints
+                }
+            }
+        }
+    }
+
     #     Describe 'Stripping common parent paths' {
 
     #         If ( (& $SafeCommands['Get-Variable'] -Name IsLinux -Scope Global -ErrorAction SilentlyContinue) -or