Set Anaconda Environment in PowerShell

AnacondaEnv.psm1

#region >> Helper Functions

function Get-FilePath {
    [CmdletBinding()]
    Param (
        [Parameter(Mandatory=$True)]
        [string]$FileNameWExtension,

        [Parameter(Mandatory=$True)]
        [string[]]$DirectoriesToSearch
    )

    [System.Collections.ArrayList]$FoundFileItems = @()
    foreach ($Dir in $DirectoriesToSearch) {
        if (Test-Path $Dir) {
            Get-ChildItem -Path $Dir -File -Recurse -Filter $FileNameWExtension | Where-Object {
                ![bool]$($_.FullName -match "\\pkgs\\")
            } | foreach {
                $null = $FoundFileItems.Add($_)
            }
        }
        else {
            Write-Warning "The directory '$Dir' does not exist!"
        }
    }

    if ($FoundFileItems.Count -eq 0) {
        Write-Error "Unable to find '$FileNameWExtension' under the following directories:`n$($DirectoriesToSearch -join "`n")`nHalting!"
        $global:FunctionResult = "1"
        return
    }
    if ($FoundFileItems.Count -eq 1) {
        $FinalFilePath = $FoundFileItems[0].FullName
    }
    if ($FoundFileItems.Count -gt 1) {
        Write-Warning "Multiple '$FileNameWExtension' files were found! Available '$FileNameWExtension' files are as follows:"
        for ($i=0; $i -lt $FoundFileItems.Count; $i++) {
            Write-Host "$i) '$($FoundFileItems[$i].FullName)'"
        }
        $ValidChoiceNumbers = 0..$($FoundFileItems.Count-1)
        $ChoiceNumber = Read-Host -Prompt "Please enter the number that corresponds to the '$FileNameWExtension' file that you would like to use"
        while ($ValidChoiceNumbers -notcontains $ChoiceNumber) {
            Write-Host "'$ChoiceNumber' is not a valid choice! Valid choices are $($ValidChoiceNumbers -join ",")"
            $ChoiceNumber = Read-Host -Prompt "Please enter the number that corresponds to the '$FileNameWExtension' file that you would like to use"
        }

        $FinalFilePath = $FoundFileItems[$ChoiceNumber].FullName
    }

    $FinalFilePath
}

function Add-ToPath {
    [CmdletBinding()]
    Param (
        [Parameter(Mandatory=$True)]
        [string]$PathToAdd,

        [Parameter(Mandatory=$False)]
        [switch]$SystemPathOnly,

        [Parameter(Mandatory=$False)]
        [switch]$PSEnvPathOnly
    )

    if (!$PSEnvPathOnly) {
        $CurrentSystemPath = $(Get-ItemProperty -Path 'Registry::HKEY_LOCAL_MACHINE\System\CurrentControlSet\Control\Session Manager\Environment' -Name PATH).Path
        $CurrentSystemPathArray = $CurrentSystemPath -split ";"
        if ($CurrentSystemPathArray -notcontains $PathToAdd) {
            $UpdatedSystemPath = "$PathToAdd;$CurrentSystemPath"
            Set-ItemProperty -Path "Registry::HKEY_LOCAL_MACHINE\System\CurrentControlSet\Control\Session Manager\Environment" -Name PATH -Value $UpdatedSystemPath
        }
    }

    if (!$SystemPathOnly) {
        $CurrentEnvPathArray = $env:Path -split ";"
        if ($CurrentEnvPathArray -notcontains $PathToAdd) {
            $env:Path = "$PathToAdd;$env:Path"
        }
    }
}

#endregion >> Helper Functions

function Set-AnacondaEnv {
    [CmdletBinding()]
    Param (
        [Parameter(Mandatory=$False)]
        [string]$AnacondaDirectoryPath,

        [Parameter(Mandatory=$False)]
        [switch]$HideOutput,

        [Parameter(Mandatory=$False)]
        [switch]$ChangeMachineEnv
    )

    #region >> Prep

    $OriginalSystemPath = $(Get-ItemProperty -Path 'Registry::HKEY_LOCAL_MACHINE\System\CurrentControlSet\Control\Session Manager\Environment' -Name PATH).Path
    $OriginalPSEnvPath = $env:Path
    $OriginalPYTHONIOENCODING = $env:PYTHONIOENCODING
    $OriginalCONDA_EXE = $env:CONDA_EXE
    $OriginalCONDA_NEW_ENV = $env:CONDA_NEW_ENV
    $OriginalCONDA_PS1_BACKUP = $env:CONDA_PS1_BACKUP


    if ($AnacondaDirectoryPath) {
        $AnacondaDirectoryName = $AnacondaDirectoryPath | Split-Path -Leaf
    }
    else {
        $AnacondaDirectoryName = "Anaconda3"
    }

    # Make sure we can find the Anaconda3 directory
    if (!$PSBoundParameters['AnacondaDirectoryPath']) {
        $PotentialDirectoriesToSearch = @(
            "C:\tools\$AnacondaDirectoryName"
            "C:\ProgramData\$AnacondaDirectoryName"
            "$HOME\$AnacondaDirectoryName"
        )

        $DirectoriesToSearch = $PotentialDirectoriesToSearch | foreach {
            if (Test-Path $_) {
                $_
            }
        }
    }
    else {
        if (!$(Test-Path $AnacondaDirectoryPath)) {
            Write-Error "The path '$AnacondaDirectoryPath' does not exist! Halting!"
            $global:FunctionResult = "1"
            return
        }

        $DirectoriesToSearch = $AnacondaDirectoryPath
    }

    # Find python.exe
    try {
        $PythonExePath = Get-FilePath -FileNameWExtension "python.exe" -DirectoriesToSearch $DirectoriesToSearch -ErrorAction Stop
        if (!$PythonExePath) {throw "The Get-FilePath function failed! Halting!"}
        $PythonParentDir = $PythonExePath | Split-Path -Parent
        $PythonCmd = $PythonExePath | Split-Path -Leaf
    }
    catch {
        Write-Error $_
        $global:FunctionResult = "1"
        return
    }

    # Add python.exe to System PATH and $env:Path
    try {
        if ($ChangeMachineEnv) {
            $null = Add-ToPath -PathToAdd $PythonParentDir -ErrorAction Stop
        }
        else {
            $null = Add-ToPath -PathToAdd $PythonParentDir -PSEnvPathOnly -ErrorAction Stop
        }

        if (![bool]$(Get-Command $PythonCmd -ErrorAction SilentlyContinue)) {
            throw "Did not successfully add '$PythonCmd' to System and PowerShell environment paths! Halting!"
        }
    }
    catch {
        Write-Error $_
        $global:FunctionResult = "1"
        return
    }

    # Set PYTHONIOENCODING PowerShell-Specific Environment Variable
    $env:PYTHONIOENCODING = python.exe -c 'import ctypes; print(ctypes.cdll.kernel32.GetACP())'
    if (!$env:PYTHONIOENCODING) {
        Write-Error "Unable to determine `$env:PYTHONIOENCODING! Halting!"
        $global:FunctionResult = "1"
        return
    }
    if ($ChangeMachineEnv) {
        # Set PYTHONIOENCODING System Environment Variable
        [Environment]::SetEnvironmentVariable("PYTHONIOENCODING", $env:PYTHONIOENCODING, "Machine")
    }

    # Set CONDA_NEW_ENV PowerShell-Specific Environment Variable
    $env:CONDA_NEW_ENV = $PythonParentDir
    if ($ChangeMachineEnv) {
        # Set CONDA_NEW_ENV System Environment Variable
        [Environment]::SetEnvironmentVariable("CONDA_NEW_ENV", $env:CONDA_NEW_ENV, "Machine")
    }

    # Find conda.exe
    try {
        $CondaExePath = Get-FilePath -FileNameWExtension "conda.exe" -DirectoriesToSearch $DirectoriesToSearch -ErrorAction Stop
        if (!$CondaExePath) {throw "The Get-FilePath function failed! Halting!"}
        $CondaParentDir = $CondaExePath | Split-Path -Parent
        $CondaCmd = $CondaExePath | Split-Path -Leaf
    }
    catch {
        Write-Error $_
        $global:FunctionResult = "1"
        return
    }

    # Add conda.exe to System PATH and $env:Path
    try {
        if ($ChangeMachineEnv) {
            $null = Add-ToPath -PathToAdd $CondaParentDir -ErrorAction Stop
        }
        else {
            $null = Add-ToPath -PathToAdd $CondaParentDir -PSEnvPathOnly -ErrorAction Stop
        }

        if (![bool]$(Get-Command $CondaCmd -ErrorAction SilentlyContinue)) {
            throw "Did not successfully add '$CondaCmd' to System and PowerShell environment paths! Halting!"
        }
    }
    catch {
        Write-Error $_
        $global:FunctionResult = "1"
        return
    }

    # Set CONDA_EXE PowerShell-Specific Environment Variable
    $env:CONDA_EXE = $CondaExePath
    if ($ChangeMachineEnv) {
        # Set CONDA_EXE System Environment Variable
        [Environment]::SetEnvironmentVariable("CONDA_EXE", $env:CONDA_EXE, "Machine")
    }

    # Set CONDA_PS1_BACKUP PowerShell-Specific Environment Variable
    $env:CONDA_PS1_BACKUP = $(Get-Location).Path
    if ($ChangeMachineEnv) {
        # Set CONDA_PS1_BACKUP System Environment Variable
        [Environment]::SetEnvironmentVariable("CONDA_PS1_BACKUP", $env:CONDA_PS1_BACKUP, "Machine")
    }

    $NewSystemPath = $(Get-ItemProperty -Path 'Registry::HKEY_LOCAL_MACHINE\System\CurrentControlSet\Control\Session Manager\Environment' -Name PATH).Path
    $NewPSEnvPath = $env:Path
    $NewPYTHONIOENCODING = $env:PYTHONIOENCODING
    $NewCONDA_EXE = $env:CONDA_EXE
    $NewCONDA_NEW_ENV = $env:CONDA_NEW_ENV
    $NewCONDA_PS1_BACKUP = $env:CONDA_PS1_BACKUP

    $FinalNewSystemPath = if ($OriginalSystemPath -ne $NewSystemPath) {$NewSystemPath} else {"NoChange"}
    $FinalPSEnvPath = if ($OriginalPSEnvPath -ne $NewPSEnvPath) {$NewPSEnvPath} else {"NoChange"}
    $FinalPyEn = if ($OriginalPYTHONIOENCODING -ne $NewPYTHONIOENCODING) {$NewPYTHONIOENCODING} else {"NoChange"}
    $FinalCondaExe = if ($OriginalCONDA_EXE -ne $NewCONDA_EXE) {$NewCONDA_EXE} else {"NoChange"}
    $FinalCondaNewEnv = if ($OriginalCONDA_NEW_ENV -ne $NewCONDA_NEW_ENV) {$NewCONDA_NEW_ENV} else {"NoChange"}
    $FInalCondaPS1 = if ($OriginalCONDA_PS1_BACKUP -ne $NewCONDA_PS1_BACKUP) {$NewCONDA_PS1_BACKUP} else {"NoChange"}

    if (!$HideOutput) {
        [pscustomobject]@{
            SystemPathChanges       = @{Original = $OriginalSystemPath; New = $FinalNewSystemPath}
            PSEnvPathChanges        = @{Original = $OriginalPSEnvPath; New = $FinalPSEnvPath}
            PYTHONIOENCODING        = @{Original = $OriginalPYTHONIOENCODING; New = $FinalPyEn}
            CONDA_EXE               = @{Original = $OriginalCONDA_EXE; New = $FinalCondaExe}
            CONDA_NEW_ENV           = @{Original = $OriginalCONDA_NEW_ENV; New = $FinalCondaNewEnv}
            CONDA_PS1_BACKUP        = @{Original = $OriginalCONDA_NEW_ENV; New = $FInalCondaPS1}
        }

        Write-Host "Environment set successfully!" -ForegroundColor Green
    }
}

function Revert-AnacondaEnv {
    [CmdletBinding()]
    Param (
        [Parameter(Mandatory=$False)]
        [pscustomobject]$SetAnacondaEnvOutput
    )

    # Validate $SetAnacondaEnvOutput object
    $NotePropertiesOutputShouldHave = @(
        "SystemPathChanges"
        "PSEnvPathChanges"
        "PYTHONIOENCODING"
        "CONDA_EXE"
        "CONDA_NEW_ENV"
        "CONDA_PS1_BACKUP"
    )

    foreach ($ValidProperty in $NotePropertiesOutputShouldHave) {
        if ($($SetAnacondaEnvOutput | Get-Member).Name -notcontains $ValidProperty) {
            Write-Error "The pscustomobject provided to the -SetAnacondaEnvOutput parameter is invalid! Halting!"
            $global:FunctionResult = "1"
            return
        }
    }

    # We're not going to halt on any errors, because we want to report all errors at the end, so
    # we need to collect success/failure as we go...
    [System.Collections.ArrayList]$SuccessfulReversions = @()
    [System.Collections.ArrayList]$FailedReversions = @()

    # Revert System Path and $env:Path
    try {
        if ($SetAnacondaEnvOutput.SystemPathChanges.New -ne "NoChange") {
            $SetItemPropertySplatParams = @{
                Path        = "Registry::HKEY_LOCAL_MACHINE\System\CurrentControlSet\Control\Session Manager\Environment"
                Name        = "PATH"
                Value       = $SetAnacondaEnvResult.SystemPathChanges.Original
            }
            Set-ItemProperty @SetItemPropertySplatParams

            $SystemPathState = [pscustomobject]@{
                EnvironmentCharacteristic   = "SystemPath"
                AttemptedValue              = $SetAnacondaEnvResult.SystemPathChanges.Original
                CurrentValue                = $(Get-ItemProperty -Path 'Registry::HKEY_LOCAL_MACHINE\System\CurrentControlSet\Control\Session Manager\Environment' -Name PATH).Path
            }
            $null = $SuccessfulReversions.Add($SystemPathState)
        }
    }
    catch {
        Write-Warning $($_ | Out-String)
        Write-Warning "Unable to revert System Path!"

        $SystemPathState = [pscustomobject]@{
            EnvironmentCharacteristic   = "SystemPath"
            AttemptedValue              = $SetAnacondaEnvResult.SystemPathChanges.Original
            CurrentValue                = $(Get-ItemProperty -Path 'Registry::HKEY_LOCAL_MACHINE\System\CurrentControlSet\Control\Session Manager\Environment' -Name PATH).Path
        }
        $null = $FailedReversions.Add($SystemPathState)
    }

    if ($SetAnacondaEnvOutput.PSEnvPathChanges.New -ne "NoChange") {
        $env:Path = $SetAnacondaEnvOutput.PSEnvPathChanges.Original

        $PSEnvPathState = [pscustomobject]@{
            EnvironmentCharacteristic   = "PSEnvPath"
            AttemptedValue              = $SetAnacondaEnvOutput.PSEnvPathChanges.Original
            CurrentValue                = $env:Path
        }
        $null = $SuccessfulReversions.Add($PSEnvPathState)
    }

    # Revert $env:PYTHONIOENCODING...
    if ($SetAnacondaEnvOutput.PYTHONIOENCODING.New -ne "NoChange") {
        try {
            $env:PYTHONIOENCODING = $SetAnacondaEnvOutput.PYTHONIOENCODING.Original
            [Environment]::SetEnvironmentVariable("PYTHONIOENCODING", $SetAnacondaEnvOutput.PYTHONIOENCODING.Original, "Machine")

            $PyIOEnState = [pscustomobject]@{
                EnvironmentCharacteristic   = "PYTHONIOENCODING"
                AttemptedValue              = $SetAnacondaEnvOutput.PYTHONIOENCODING.Original
                CurrentValue                = $SetAnacondaEnvOutput.PYTHONIOENCODING.Original
            }
            $null = $SuccessfulReversions.Add($PyIOEnState)
        }
        catch {
            Write-Warning $($_ | Out-String)
            Write-Warning "Unable to revert the System Environment Variable 'PYTHONIOENCODING'!"

            $PyIOEnState = [pscustomobject]@{
                EnvironmentCharacteristic   = "PYTHONIOENCODING"
                AttemptedValue              = $SetAnacondaEnvOutput.PYTHONIOENCODING.Original
                CurrentValue                = $SetAnacondaEnvOutput.PYTHONIOENCODING.New
            }
            $null = $FailedReversions.Add($PyIOEnState)
        }
    }

    if ($SetAnacondaEnvOutput.CONDA_EXE.New -ne "NoChange") {
        try {
            $env:CONDA_EXE = $SetAnacondaEnvOutput.CONDA_EXE.Original
            [Environment]::SetEnvironmentVariable("CONDA_EXE", $SetAnacondaEnvOutput.CONDA_EXE.Original, "Machine")

            $CondaExeState = [pscustomobject]@{
                EnvironmentCharacteristic   = "CONDA_EXE"
                AttemptedValue              = $SetAnacondaEnvOutput.CONDA_EXE.Original
                CurrentValue                = $SetAnacondaEnvOutput.CONDA_EXE.Original
            }
            $null = $SuccessfulReversions.Add($CondaExeState)
        }
        catch {
            Write-Warning $($_ | Out-String)
            Write-Warning "Unable to revert the System Environment Variable 'CONDA_EXE'!"

            $CondaExeState = [pscustomobject]@{
                EnvironmentCharacteristic   = "CONDA_EXE"
                AttemptedValue              = $SetAnacondaEnvOutput.CONDA_EXE.Original
                CurrentValue                = $SetAnacondaEnvOutput.CONDA_EXE.New
            }
            $null = $FailedReversions.Add($CondaExeState)
        }
    }

    if ($SetAnacondaEnvOutput.CONDA_NEW_ENV.New -ne "NoChange") {
        try {
            $env:CONDA_NEW_ENV = $SetAnacondaEnvOutput.CONDA_NEW_ENV.Original
            [Environment]::SetEnvironmentVariable("CONDA_NEW_ENV", $SetAnacondaEnvOutput.CONDA_NEW_ENV.Original, "Machine")

            $CondaNewEnvState = [pscustomobject]@{
                EnvironmentCharacteristic   = "CONDA_NEW_ENV"
                AttemptedValue              = $SetAnacondaEnvOutput.CONDA_NEW_ENV.Original
                CurrentValue                = $SetAnacondaEnvOutput.CONDA_NEW_ENV.Original
            }
            $null = $SuccessfulReversions.Add($CondaNewEnvState)
        }
        catch {
            Write-Warning $($_ | Out-String)
            Write-Warning "Unable to revert the System Environment Variable 'CONDA_NEW_ENV'!"

            $CondaNewEnvState = [pscustomobject]@{
                EnvironmentCharacteristic   = "CONDA_NEW_ENV"
                AttemptedValue              = $SetAnacondaEnvOutput.CONDA_NEW_ENV.Original
                CurrentValue                = $SetAnacondaEnvOutput.CONDA_NEW_ENV.New
            }
            $null = $FailedReversions.Add($CondaNewEnvState)
        }
    }

    if ($SetAnacondaEnvOutput.CONDA_PS1_BACKUP.New -ne "NoChange") {
        try {
            $env:CONDA_PS1_BACKUP = $SetAnacondaEnvOutput.CONDA_PS1_BACKUP.Original
            [Environment]::SetEnvironmentVariable("CONDA_PS1_BACKUP", $SetAnacondaEnvOutput.CONDA_PS1_BACKUP.Original, "Machine")

            $CondaPS1State = [pscustomobject]@{
                EnvironmentCharacteristic   = "CONDA_PS1_BACKUP"
                AttemptedValue              = $SetAnacondaEnvOutput.CONDA_PS1_BACKUP.Original
                CurrentValue                = $SetAnacondaEnvOutput.CONDA_PS1_BACKUP.Original
            }
            $null = $SuccessfulReversions.Add($CondaPS1State)
        }
        catch {
            Write-Warning $($_ | Out-String)
            Write-Warning "Unable to revert the System Environment Variable 'CONDA_PS1_BACKUP'!"

            $CondaPS1State = [pscustomobject]@{
                EnvironmentCharacteristic   = "CONDA_PS1_BACKUP"
                AttemptedValue              = $SetAnacondaEnvOutput.CONDA_PS1_BACKUP.Original
                CurrentValue                = $SetAnacondaEnvOutput.CONDA_PS1_BACKUP.New
            }
            $null = $FailedReversions.Add($CondaPS1State)
        }
    }

    # Prep Output
    $Output = @{}

    if ($SuccessfulReversions.Count -gt 0) {
        $Output.Add("SuccessfulReversions",$SuccessfulReversions)
    }
    if ($FailedReversions.Count -gt 0) {
        $Output.Add("FailedReversions",$FailedReversions)
    }

    if ($SuccessfulReversions.Count -eq 0 -and $FailedReversions.Count -eq 0) {
        Write-Host "Nothing was reverted because nothing was changed by the Set-AnacondaEnv function in the first place!" -ForegroundColor Green
    }
    if ($SuccessfulReversions.Count -gt 0 -and $FailedReversions.Count -eq 0) {
        Write-Host "Reverted environment successfully!" -ForegroundColor Green
    }
    
    $Output
}