Update PowerShell type data only temporary, while a given script block runs

TemporaryTypeData.psm1

<# Copyright 2022 zett42
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files 
(the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, 
publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so,
subject to the following conditions:
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES 
OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 
LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. #>

$ErrorActionPreference = 'Stop'

function Invoke-WithTemporaryTypeData {
    <#
    .SYNOPSIS
        Modify PowerShell type data temporary.
    .DESCRIPTION
        Modify type data by calling Update-TypeData, run a script block and restore type data afterwards.           

    .PARAMETER UpdateTypeDataArgs
        Arguments for Update-TypeData. By passing an array, multiple types can be updated.
    .PARAMETER NewScope
        Run the given script block in a new scope instead of in the caller's scope.
        NOTE: For running the script block in the caller's scope, this function must be contained in a module.
    .PARAMETER ScriptBlock
        The script block to run while the new type data is in effect.
    #>

    [CmdletBinding()]
    param (
        [Parameter(Mandatory)]
        [Hashtable[]] $UpdateTypeDataArgs,

        [switch] $NewScope,

        [Parameter(Mandatory)]
        [scriptblock] $ScriptBlock
    )
    
    # Get unique type names as multiple updates to the same type could have been requested.
    $uniqueTypeNames = $UpdateTypeDataArgs.TypeName | Sort-Object -Unique

    # Save preexisting type data, if any.
    $prevTypeData = Get-TypeData $uniqueTypeNames

    try {
        # Modify the type data.
        foreach( $utdArgs in $UpdateTypeDataArgs ) { 
            Update-TypeData @utdArgs -Force 
        }

        # Run the script block either in a new scope or in the caller's scope.
        if( $NewScope ) { & $ScriptBlock } else { . $ScriptBlock }
    }
    finally {
        # Restore the original behavior:
        # Remove the override again...
        Get-TypeData $uniqueTypeNames | Remove-TypeData

        # ... and restore the previous data, if any.
        if( $prevTypeData ) {
            Update-TypeData -TypeData $prevTypeData
        }
    }
}

TemporaryTypeData.Tests.ps1

# Pester tests for TemporaryTypeData.psm1

BeforeAll {
    Import-Module $PSScriptRoot\TemporaryTypeData.psm1 -Force
}

Describe 'Invoke-WithTemporaryTypeData' {

    It 'invokes a script block with temporary type data' {

        # Arguments for Update-TypeData.
        $typeDataArgs = @{
            TypeName   = 'System.Boolean' 
            MemberType = 'ScriptMethod' 
            MemberName = 'ToString'
            Value      = { if ($this) { 'yes' } else { 'no' } }
        }, @{
            TypeName   = 'System.Byte' 
            MemberType = 'ScriptMethod' 
            MemberName = 'ToString'
            Value      = { "[$this]" }
        }, @{
            TypeName   = 'System.Byte' 
            MemberType = 'ScriptMethod' 
            MemberName = 'TimesTwo'
            Value      = { $this * 2 }
        }

        $scriptBlockInvoked = $false

        Invoke-WithTemporaryTypeData $typeDataArgs {

            # This should print using the ToString() overrides.
            $true     | Out-String -NoNewline | Should -Be 'yes'
            $false    | Out-String -NoNewline | Should -Be 'no'
            [byte] 42 | Out-String -NoNewline | Should -Be '[42]'
         
            ([byte] 21).TimesTwo() | Should -Be 42 

            $scriptBlockInvoked = $true
        }

        # Test if script block has been invoked in caller's scope.
        $scriptBlockInvoked | Should -BeTrue

        # This should print regularly again.
        $true     | Out-String -NoNewline | Should -Be 'True'
        $false    | Out-String -NoNewline | Should -Be 'False'
        [byte] 42 | Out-String -NoNewline | Should -Be '42'

        # This should throw as the method TimesTwo() no longer exists for the [byte] type.
        { ([byte] 21).TimesTwo() } | Should -Throw
    }

    It 'supports nested function calls' {

        # Arguments for Update-TypeData.
        $typeDataArgs = @{
            TypeName   = 'System.Boolean' 
            MemberType = 'ScriptMethod' 
            MemberName = 'ToString'
            Value      = { if ($this) { 'yes' } else { 'no' } }
        }, @{
            TypeName   = 'System.Byte' 
            MemberType = 'ScriptMethod' 
            MemberName = 'ToString'
            Value      = { "[$this]" }
        }

        $scriptBlockInvoked = $false

        Invoke-WithTemporaryTypeData $typeDataArgs {

            # This should print using the ToString() overrides.
            $true     | Out-String -NoNewline | Should -Be 'yes'
            $false    | Out-String -NoNewline | Should -Be 'no'
            [byte] 42 | Out-String -NoNewline | Should -Be '[42]'

            $typeDataArgs2 = @{
                TypeName   = 'System.Boolean' 
                MemberType = 'ScriptMethod' 
                MemberName = 'ToString'
                Value      = { if ($this) { '1' } else { '0' } }
            }, @{
                TypeName   = 'System.Byte' 
                MemberType = 'ScriptMethod' 
                MemberName = 'ToString'
                Value      = { "#$this" }
            }
    
            Invoke-WithTemporaryTypeData $typeDataArgs2 {

                # This should print using the 2nd level ToString() overrides.
                $true     | Out-String -NoNewline | Should -Be '1'
                $false    | Out-String -NoNewline | Should -Be '0'
                [byte] 42 | Out-String -NoNewline | Should -Be '#42'
                
                $scriptBlockInvoked = $true               
            }

            # This should print using the 1st level ToString() overrides again.
            $true     | Out-String -NoNewline | Should -Be 'yes'
            $false    | Out-String -NoNewline | Should -Be 'no'
            [byte] 42 | Out-String -NoNewline | Should -Be '[42]'
        }

        $scriptBlockInvoked | Should -BeTrue

        # This should print regularly again.
        $true     | Out-String -NoNewline | Should -Be 'True'
        $false    | Out-String -NoNewline | Should -Be 'False'
        [byte] 42 | Out-String -NoNewline | Should -Be '42'
    }

    It 'invokes a script block in new scope' {

        # Arguments for Update-TypeData.
        $typeDataArgs = @{
            TypeName   = 'System.Byte' 
            MemberType = 'ScriptMethod' 
            MemberName = 'ToString'
            Value      = { "[$this]" }
        }

        $scriptBlockInvoked = $false

        Invoke-WithTemporaryTypeData $typeDataArgs -NewScope {

            # This should print using the ToString() overrides.
            [byte] 42 | Out-String -NoNewline | Should -Be '[42]'
         
            # This succeeds to set $scriptBlockInvoked = $true only because of -NewScope argument.  
            Set-Variable scriptBlockInvoked $true -Scope 1
        }

        # Test if script block has been invoked in child scope.
        $scriptBlockInvoked | Should -BeTrue

        # This should print regularly again.
        [byte] 42 | Out-String -NoNewline | Should -Be '42'
    }    

    It 'restores type data even when an exception gets thrown' {

        # Arguments for Update-TypeData.
        $typeDataArgs = @{
            TypeName   = 'System.Byte' 
            MemberType = 'ScriptMethod' 
            MemberName = 'ToString'
            Value      = { "[$this]" }
        }

        $scriptBlockInvoked = $false

        try {
            Invoke-WithTemporaryTypeData $typeDataArgs {

                # This should print using the ToString() overrides.
                [byte] 42 | Out-String -NoNewline | Should -Be '[42]'

                $scriptBlockInvoked = $true

                throw 'just a test'
            }
        }
        catch {
            #Write-Debug $_ -Debug
        }

        $scriptBlockInvoked | Should -BeTrue

        # This should print regularly again.
        [byte] 42 | Out-String -NoNewline | Should -Be '42'
    }
}