improve exceptional behavior of sh.cd

lib/sh/cd.py

@@ -1,7 +1,8 @@
 #!/usr/bin/env py.test
 from __future__ import annotations
 
-import contextlib
+from contextlib import contextmanager
+from os import chdir
 from os import environ
 
 from lib import json as JSON
@@ -19,39 +20,39 @@
 Command = tuple[object, ...]
 
 
-@contextlib.contextmanager
+@contextmanager
 def cd(
     dirname: OSPath, env: Environ = environ, *, direnv: bool = True
 ) -> Generator[OSPath]:
     oldpwd = OSPath(env["PWD"])
 
-    # double-check that $PWD stays accurate:
-    assert oldpwd.samefile(OSPath.cwd()), (oldpwd, OSPath.cwd())
-
     newpwd = oldpwd / dirname
-    if newpwd.samefile(oldpwd):  # we're already there
+    cwd = OSPath.cwd()
+    if newpwd == oldpwd and cwd.samefile(newpwd):  # we're already there
         yield oldpwd
         return
 
-    env["PWD"] = str(newpwd)
     xtrace(("cd", dirname))
-    with contextlib.chdir(dirname):
-        if direnv:
-            run(("direnv", "allow"))
-            direnv_json: JSON.Value = json(("direnv", "export", "json"))
-            if direnv_json is None:
-                pass  # nothing to do
-            elif isinstance(direnv_json, dict):
-                for key, value in direnv_json.items():
-                    if value is None:
-                        env.pop(key, None)
-                    else:
-                        assert isinstance(value, str), value
-                        env[key] = value
-            else:
-                raise AssertionError(f"expected dict, got {type(direnv_json)}")
-        yield dirname
-
-        # show the un-cd, and reflect it in $PWD
+    env["PWD"] = str(newpwd)
+    chdir(dirname)
+    if direnv:
+        run(("direnv", "allow"))
+        direnv_json: JSON.Value = json(("direnv", "export", "json"))
+        if direnv_json is None:
+            pass  # nothing to do
+        elif isinstance(direnv_json, dict):
+            for key, value in direnv_json.items():
+                if value is None:
+                    env.pop(key, None)
+                else:
+                    assert isinstance(value, str), value
+                    env[key] = value
+        else:
+            raise AssertionError(f"expected dict, got {type(direnv_json)}")
+    try:
+        yield newpwd
+    finally:  # undo the cd and log it
+        chdir(oldpwd)
         env["PWD"] = str(oldpwd)
-        xtrace(("cd", oldpwd))
+        xtrace(("popd",))
+        info(oldpwd, "<-", newpwd)