Merge pull request #23 from hsaunders1904/make_pwsh_activator_more_robust

Make pwsh activator more robust

Author: hsaunders1904

PyAutoEnv.ps1

@@ -13,11 +13,7 @@
 #
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
-$PyAutoEnvDir = "${PSScriptRoot}"
-
-if (Test-Path alias:cd) {
-  Remove-Alias -Name cd -Force -Scope Global
-}
+$pyAutoEnvDir = "${PSScriptRoot}"
 
 <#
 .SYNOPSIS
@@ -35,12 +31,12 @@ function Invoke-PyAutoEnv() {
   if (-Not (Get-Command "python" -ErrorAction SilentlyContinue)) {
     return
   }
-  $PyAutoEnv = Join-Path "${PyAutoEnvDir}" "pyautoenv.py"
-  if (Test-Path "${PyAutoEnv}") {
-    $Expression = "$(python "${PyAutoEnv}" --pwsh)"
-    if (${Expression}) {
-      Invoke-Expression "${Expression}"
-     }
+  $pyAutoEnv = Join-Path "${pyAutoEnvDir}" "pyautoenv.py"
+  if (Test-Path "${pyAutoEnv}") {
+    $expression = "$(python "${pyAutoEnv}" --pwsh)"
+    if (${expression}) {
+      Invoke-Expression "${expression}"
+    }
   }
 }
 
@@ -51,12 +47,46 @@ function Invoke-PyAutoEnv() {
   https://github.com/hsaunders1904/pyautoenv/
 #>
 function Invoke-PyAutoEnvVersion() {
-  $PyAutoEnv = Join-Path "${PyAutoEnvDir}" "pyautoenv.py"
-  python "${PyAutoEnv}" --version
+  $pyAutoEnv = Join-Path "${pyAutoEnvDir}" "pyautoenv.py"
+  python "${pyAutoEnv}" --version
 }
 
-function cd() {
-  Set-Location @Args && Invoke-PyAutoEnv
+<#
+.SYNOPSIS
+  Create a proxy function definition for a Cmdlet that executes pyautoenv.
+.LINK
+  https://github.com/hsaunders1904/pyautoenv/
+#>
+function New-PyAutoEnvProxyFunctionDefinition([string] $commandString)
+{
+  # Generate base code for the Proxy function.
+  $originalCommand = Get-Command -Name "$commandString" -CommandType Cmdlet
+  $metaData = New-Object System.Management.Automation.CommandMetaData $originalCommand
+  $proxyCode = [System.Management.Automation.ProxyCommand]::Create($metaData)
+
+  # Find the 'end' block of Set-Location's source.
+  $ast = [System.Management.Automation.Language.Parser]::ParseInput($proxyCode, [ref]$null, [ref]$null)
+  $endBlock = $ast.EndBlock.Extent.Text
+  $endBlockClosingIndex = $endBlock.LastIndexOf('}')
+  if ($endBlockClosingIndex -Le 0) {
+    # If we can't find the opening brace, something's not right, so exit early
+    # without editing the proxy to avoid breaking things.
+    $body = $ast.ToString()
+    return "function $commandString {`n${body}`n}"
+  }
+
+  # Insert the pyautoenv function call into the 'end' block of the proxy code.
+  $tab = "    "
+  $insert = "`n${tab}try {`n${tab}${tab}Invoke-PyAutoEnv`n${tab}} catch {}`n"
+  $newEndBlockOpen = $endBlock.Substring(0, $endBlockClosingIndex) + $insert
+  $newEndBlock = $newEndBlockOpen + $endBlock.Substring($endBlockClosingIndex)
+  $updatedProxyCmd = $proxyCode.Replace($endBlock, $newEndBlock)
+  return "function global:$commandString {`n$updatedProxyCmd`n}"
 }
 
-Invoke-PyAutoEnv
+foreach ($commandName in ("Set-Location", "Push-Location", "Pop-Location")) {
+  Invoke-Expression (& {
+    (New-PyAutoEnvProxyFunctionDefinition "$commandName" | Out-String)
+  })
+}
+Invoke-PyAutoEnv  # Look for environment in initial directory.

README.md

@@ -95,9 +95,6 @@ To enable the application in PowerShell, dot the `.ps1` file.
 
 Add this to your profile to activate the application permanently.
 
-Note that this application hooks into and re-aliases `cd`.
-Therefore `pyautoenv` does not support the use of `Set-Location`.
-
 </details>
 
 ## Options

pyautoenv.py

@@ -373,26 +373,27 @@ def poetry_project_name(directory: str) -> Union[str, None]:
 
 def activator(env_directory: str, args: Args) -> str:
     """Get the activator script for the environment in the given directory."""
-    dir_name = "Scripts" if operating_system() == Os.WINDOWS else "bin"
+    is_windows = operating_system() == Os.WINDOWS
+    dir_name = "Scripts" if is_windows else "bin"
     if args.fish:
         script = "activate.fish"
     elif args.pwsh:
-        poetry_dir = poetry_cache_dir()
-        if (
-            poetry_dir is not None
-            and env_directory.startswith(poetry_dir)
-            and operating_system() != Os.WINDOWS
-        ):
-            # In poetry environments on *NIX systems, this activator has
-            # a lowercase A.
-            script = "activate.ps1"
+        if is_windows:
+            script = "Activate.ps1"
         else:
-            # In venv environments, and Windows poetry environments,
-            # this activator has an uppercase A.
+            # PowerShell activation scripts on Nix systems have some
+            # slightly inconsistent naming. When using Poetry or uv, the
+            # activation script is lower case, using the venv module,
+            # the script is title case.
+            # We can't really know what was used to generate the venv
+            # so just check which activation script exists.
+            script_path = os.path.join(env_directory, dir_name, "activate.ps1")
+            if os.path.isfile(script_path):
+                return script_path
             script = "Activate.ps1"
     else:
         script = "activate"
-    return os.path.join(env_directory, dir_name, f"{script}")
+    return os.path.join(env_directory, dir_name, script)
 
 
 @lru_cache(maxsize=128)

tests/test_venv.py

@@ -63,7 +63,7 @@ def fs(self, fs: FakeFilesystem) -> FakeFilesystem:
         fs.create_dir(self.PY_PROJ / "src")
         fs.create_file(self.VENV_DIR / "bin" / "activate")
         fs.create_file(self.VENV_DIR / "bin" / "activate.fish")
-        fs.create_file(self.VENV_DIR / "bin" / "Activate.ps1")
+        fs.create_file(self.VENV_DIR / "bin" / "activate.ps1")
         fs.create_file(self.VENV_DIR / "Scripts" / "activate")
         fs.create_file(self.VENV_DIR / "Scripts" / "Activate.ps1")
         fs.create_dir("not_a_venv")
@@ -200,7 +200,7 @@ class TestVenvBashLinux(VenvTester):
 
 
 class TestVenvPwshLinux(VenvTester):
-    activator = "bin/Activate.ps1"
+    activator = "bin/activate.ps1"
     flag = "--pwsh"
     os = pyautoenv.Os.LINUX