PSPlumbingRules

0.1.0

Rules/PSAvoidOutputTypeDrift.psm1

Set-StrictMode -Version Latest

$script:RuleName = 'PSAvoidOutputTypeDrift'
$script:RuleMessage = 'Avoid changing emitted output type across switch-driven branches in pipeline-facing commands. Keep output shape stable so the command remains predictable in pipelines.'

function Measure-PSAvoidOutputTypeDrift {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.ScriptBlockAst] $ScriptBlockAst
    )

    $functions = $ScriptBlockAst.FindAll({
        param($Ast)
        $Ast -is [System.Management.Automation.Language.FunctionDefinitionAst]
    }, $true)

    foreach ($functionAst in $functions) {
        if (-not (Test-IsPipelineFacingCommand -FunctionAst $functionAst)) {
            continue
        }

        if (-not (Test-HasOutputTypeDrift -FunctionAst $functionAst)) {
            continue
        }

        [Microsoft.Windows.PowerShell.ScriptAnalyzer.Generic.DiagnosticRecord]::new(
            $script:RuleMessage,
            $functionAst.Extent,
            $script:RuleName,
            [Microsoft.Windows.PowerShell.ScriptAnalyzer.Generic.DiagnosticSeverity]::Warning,
            $null,
            $null,
            $null
        )
    }
}

function Test-IsPipelineFacingCommand {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.FunctionDefinitionAst] $FunctionAst
    )

    if (-not $FunctionAst.Body.ParamBlock) {
        return $false
    }

    $hasCmdletBinding = $FunctionAst.Body.ParamBlock.Attributes | Where-Object {
        $_.TypeName.GetReflectionType().FullName -eq 'System.Management.Automation.CmdletBindingAttribute'
    }

    if (-not $hasCmdletBinding) {
        return $false
    }

    return $FunctionAst.Body.ProcessBlock -ne $null
}

function Test-HasOutputTypeDrift {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.FunctionDefinitionAst] $FunctionAst
    )

    $switchParameterNames = @(Get-SwitchParameterNames -FunctionAst $FunctionAst)
    if ($switchParameterNames.Count -eq 0) {
        return $false
    }

    $ifStatements = Get-OutputBearingIfStatements -FunctionAst $FunctionAst
    foreach ($ifStatement in $ifStatements) {
        if (-not (Test-IfStatementUsesSwitchCondition -IfStatementAst $ifStatement -SwitchParameterNames $switchParameterNames)) {
            continue
        }

        $branchKinds = @(Get-IfStatementOutputKinds -IfStatementAst $ifStatement)
        if ($branchKinds.Count -gt 1) {
            return $true
        }
    }

    return $false
}

function Get-SwitchParameterNames {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.FunctionDefinitionAst] $FunctionAst
    )

    @(
        $FunctionAst.Body.ParamBlock.Parameters |
            Where-Object { $_.StaticType -eq [System.Management.Automation.SwitchParameter] } |
            ForEach-Object { $_.Name.VariablePath.UserPath }
    )
}

function Get-OutputBearingIfStatements {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.FunctionDefinitionAst] $FunctionAst
    )

    @(
        $FunctionAst.FindAll({
            param($Ast)
            $Ast -is [System.Management.Automation.Language.IfStatementAst]
        }, $true)
    )
}

function Test-IfStatementUsesSwitchCondition {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.IfStatementAst] $IfStatementAst,

        [Parameter(Mandatory)]
        [string[]] $SwitchParameterNames
    )

    foreach ($clause in $IfStatementAst.Clauses) {
        $variables = $clause.Item1.FindAll({
            param($Ast)
            $Ast -is [System.Management.Automation.Language.VariableExpressionAst]
        }, $true)

        foreach ($variable in $variables) {
            if ($SwitchParameterNames -contains $variable.VariablePath.UserPath) {
                return $true
            }
        }
    }

    return $false
}

function Get-IfStatementOutputKinds {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.IfStatementAst] $IfStatementAst
    )

    $kinds = New-Object System.Collections.Generic.List[string]

    foreach ($clause in $IfStatementAst.Clauses) {
        foreach ($kind in @(Get-BranchOutputKinds -StatementBlockAst $clause.Item2)) {
            if (-not $kinds.Contains($kind)) {
                $kinds.Add($kind)
            }
        }
    }

    if ($IfStatementAst.ElseClause) {
        foreach ($kind in @(Get-BranchOutputKinds -StatementBlockAst $IfStatementAst.ElseClause)) {
            if (-not $kinds.Contains($kind)) {
                $kinds.Add($kind)
            }
        }
    }
    else {
        foreach ($kind in @(Get-ContinuationOutputKinds -IfStatementAst $IfStatementAst)) {
            if (-not $kinds.Contains($kind)) {
                $kinds.Add($kind)
            }
        }
    }

    $kinds
}

function Get-BranchOutputKinds {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.Ast] $StatementBlockAst
    )

    $kinds = New-Object System.Collections.Generic.List[string]

    foreach ($statement in $StatementBlockAst.Statements) {
        $kind = Get-OutputKind -StatementAst $statement
        if ($null -eq $kind) {
            continue
        }

        if (-not $kinds.Contains($kind)) {
            $kinds.Add($kind)
        }
    }

    $kinds
}

function Get-ContinuationOutputKinds {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.IfStatementAst] $IfStatementAst
    )

    $parentBlock = $IfStatementAst.Parent
    if ($null -eq $parentBlock -or -not ($parentBlock.PSObject.Properties.Name -contains 'Statements')) {
        return @()
    }

    $kinds = New-Object System.Collections.Generic.List[string]
    $statementIndex = -1

    for ($i = 0; $i -lt $parentBlock.Statements.Count; $i++) {
        if ($parentBlock.Statements[$i] -eq $IfStatementAst) {
            $statementIndex = $i
            break
        }
    }

    if ($statementIndex -lt 0) {
        return @()
    }

    for ($i = $statementIndex + 1; $i -lt $parentBlock.Statements.Count; $i++) {
        $kind = Get-OutputKind -StatementAst $parentBlock.Statements[$i]
        if ($null -eq $kind) {
            continue
        }

        if (-not $kinds.Contains($kind)) {
            $kinds.Add($kind)
        }
    }

    $kinds
}

function Get-OutputKind {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.StatementAst] $StatementAst
    )

    $pipelineAst = $null
    if ($StatementAst -is [System.Management.Automation.Language.PipelineAst]) {
        $pipelineAst = $StatementAst
    }
    elseif ($StatementAst -is [System.Management.Automation.Language.ReturnStatementAst] -and $StatementAst.Pipeline) {
        $pipelineAst = $StatementAst.Pipeline
    }

    if ($null -eq $pipelineAst) {
        return $null
    }

    $element = $pipelineAst.PipelineElements[0]
    if ($element -is [System.Management.Automation.Language.CommandAst]) {
        return 'Command'
    }

    if ($element -isnot [System.Management.Automation.Language.CommandExpressionAst]) {
        return 'Other'
    }

    $expression = $element.Expression
    if ($expression -is [System.Management.Automation.Language.ConvertExpressionAst] -and
        $expression.Type.TypeName.FullName -eq 'pscustomobject') {
        return 'Object'
    }

    if ($expression -is [System.Management.Automation.Language.StringConstantExpressionAst] -or
        $expression -is [System.Management.Automation.Language.ExpandableStringExpressionAst]) {
        return 'String'
    }

    if ($expression -is [System.Management.Automation.Language.ConstantExpressionAst]) {
        return 'Scalar'
    }

    return 'Other'
}

Export-ModuleMember -Function Measure-PSAvoidOutputTypeDrift