Make Session.chdir() a context manager

nox/sessions.py

@@ -21,6 +21,7 @@
 import sys
 import unicodedata
 import warnings
+from types import TracebackType
 from typing import (
     Any,
     Callable,
@@ -31,6 +32,7 @@
     Optional,
     Sequence,
     Tuple,
+    Type,
     Union,
 )
 
@@ -116,6 +118,23 @@ class Status(enum.Enum):
     SKIPPED = 2
 
 
+class _WorkingDirContext:
+    def __init__(self, dir: Union[str, os.PathLike]) -> None:
+        self._prev_working_dir = os.getcwd()
+        os.chdir(dir)
+
+    def __enter__(self) -> "_WorkingDirContext":
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        os.chdir(self._prev_working_dir)
+
+
 class Session:
     """The Session object is passed into each user-defined session function.
 
@@ -210,10 +229,21 @@ def invoked_from(self) -> str:
         """
         return self._runner.global_config.invoked_from
 
-    def chdir(self, dir: Union[str, os.PathLike]) -> None:
-        """Change the current working directory."""
+    def chdir(self, dir: Union[str, os.PathLike]) -> _WorkingDirContext:
+        """Change the current working directory.
+
+        Can be used as a context manager to automatically restore the working directory::
+
+            with session.chdir("somewhere/deep/in/monorepo"):
+                # Runs in "/somewhere/deep/in/monorepo"
+                session.run("pytest")
+
+            # Runs in original working directory
+            session.run("flake8")
+
+        """
         self.log(f"cd {dir}")
-        os.chdir(dir)
+        return _WorkingDirContext(dir)
 
     cd = chdir
     """An alias for :meth:`chdir`."""

tests/test_sessions.py

@@ -164,6 +164,19 @@ def test_chdir(self, tmpdir):
         assert os.getcwd() == cdto
         os.chdir(current_cwd)
 
+    def test_chdir_ctx(self, tmpdir):
+        cdto = str(tmpdir.join("cdbby").ensure(dir=True))
+        current_cwd = os.getcwd()
+
+        session, _ = self.make_session_and_runner()
+
+        with session.chdir(cdto):
+            assert os.getcwd() == cdto
+
+        assert os.getcwd() == current_cwd
+
+        os.chdir(current_cwd)
+
     def test_invoked_from(self, tmpdir):
         cdto = str(tmpdir.join("cdbby").ensure(dir=True))
         current_cwd = os.getcwd()