Python OS Module Practice Problems & Exercises

Practice: OS Module

11 problems4 Easy4 Medium3 Hard⏱ 40-55 min

← Back to lesson

---

Easy

#1os.path Basics: Split and JoinEasy

os.pathbasenamedirnamejoin

Predict the output. Three os.path functions extract and reconstruct parts of a file path.

Python

import os

path = "/data/exports/report.csv"

filename = os.path.basename(path)
print(filename)

directory = os.path.dirname(path)
print(directory)

rebuilt = os.path.join(directory, filename)
print(rebuilt)

Solution

import os

path = "/data/exports/report.csv"

filename = os.path.basename(path)
print(filename)

directory = os.path.dirname(path)
print(directory)

rebuilt = os.path.join(directory, filename)
print(rebuilt)

Output:

report.csv
/data/exports
/data/exports/report.csv

How it works: os.path.basename() extracts the final component after the last separator — here, report.csv. os.path.dirname() returns everything before the last separator — /data/exports. os.path.join() recombines them with the correct OS separator, producing the original path.

Key insight: These three functions are the foundation of classic path manipulation. os.path.split(path) returns both dirname and basename as a tuple in one call: ('/data/exports', 'report.csv'). In modern code, pathlib.Path provides the same operations via .name, .parent, and / operator, but os.path is still dominant in legacy codebases.

Expected Output

report.csv\n/data/exports\n/data/exports/report.csv

Hints

Hint 1: os.path.basename() returns the final component of a path.

Hint 2: os.path.dirname() returns everything except the final component.

Hint 3: os.path.join() combines path components with the correct separator.

#2splitext: Separating Name from ExtensionEasy

os.pathsplitextextensions

Predict the output. os.path.splitext() splits a filename at the last dot. Watch what happens with double extensions.

Python

import os

name = "archive.tar.gz"

root, ext = os.path.splitext(name)
print((root[-4:], ext))

# Getting the full double extension requires two calls
root2, ext2 = os.path.splitext(root)
full_ext = ext2 + ext
print((root2, full_ext))

# Compare with pathlib behavior
from pathlib import Path
p = Path(name)
print(p.suffixes == [".tar", ".gz"])

Solution

import os

name = "archive.tar.gz"

root, ext = os.path.splitext(name)
print((root[-4:], ext))

root2, ext2 = os.path.splitext(root)
full_ext = ext2 + ext
print((root2, full_ext))

from pathlib import Path
p = Path(name)
print(p.suffixes == [".tar", ".gz"])

Output:

('.tar', '.gz')
('archive', '.tar.gz')
True

How it works: os.path.splitext("archive.tar.gz") splits at the last dot, returning ("archive.tar", ".gz"). We print the last 4 characters of root (.tar) and the extension (.gz). To get the full double extension .tar.gz, you must call splitext again on the root, getting ("archive", ".tar"), then concatenate: .tar + .gz = .tar.gz.

pathlib.Path.suffixes handles this more elegantly — it returns a list of all suffixes: [".tar", ".gz"].

Key insight: os.path.splitext() only knows about the last dot. This is a common source of bugs when processing .tar.gz, .tar.bz2, or .backup.sql files. If you need to handle compound extensions, either call splitext twice or use pathlib's .suffixes property.

Expected Output

('.tar', '.gz')\n('archive', '.tar.gz')\nTrue

Hints

Hint 1: os.path.splitext() splits at the LAST dot, not the first.

Hint 2: For double extensions like .tar.gz, splitext only peels off the last one.

#3os.getcwd and os.path.abspathEasy

os.getcwdabspathrealpath

Predict the output. Explore how os.getcwd(), os.path.abspath(), and os.path.isabs() relate to each other.

Python

import os

cwd = os.getcwd()

# Is getcwd always absolute?
print(os.path.isabs(cwd))

# Does abspath(".") give us the cwd?
print(os.path.abspath(".") == cwd)

# Does abspath resolve relative paths against cwd?
relative = "subdir/file.txt"
absolute = os.path.abspath(relative)
expected = os.path.join(cwd, relative)
print(absolute == expected)

Solution

import os

cwd = os.getcwd()

print(os.path.isabs(cwd))

print(os.path.abspath(".") == cwd)

relative = "subdir/file.txt"
absolute = os.path.abspath(relative)
expected = os.path.join(cwd, relative)
print(absolute == expected)

Output:

True
True
True

How it works:

1. os.getcwd() always returns an absolute path — it calls the operating system's getcwd() system call, which returns the full path from the root.

2. os.path.abspath(".") resolves . (current directory) against the working directory, producing the same result as os.getcwd().

3. os.path.abspath("subdir/file.txt") prepends the current working directory to the relative path, which is exactly what os.path.join(cwd, "subdir/file.txt") does.

Key insight: os.path.abspath() does not check whether the path actually exists — it just resolves ., .., and relative paths against the current working directory using string manipulation. If you need to resolve symlinks to the true physical path, use os.path.realpath() instead.

Expected Output

True\nTrue\nTrue

Hints

Hint 1: os.path.abspath() resolves a relative path against the current working directory.

Hint 2: os.getcwd() returns the current working directory as an absolute string.

#4os.environ: Read and DefaultEasy

os.environenvironment-variablesget

Predict the output. Safely read environment variables using os.environ.get() and handle missing keys.

Python

import os

# HOME (or USERPROFILE on Windows) almost always exists
home = os.environ.get("HOME") or os.environ.get("USERPROFILE")
if home:
    print("found_home")
else:
    print("no_home")

# A variable that almost certainly does not exist
missing = os.environ.get("XYZZZY_NOT_REAL_12345", "default_used")
print(missing)

# Direct access raises KeyError for missing keys
try:
    val = os.environ["XYZZZY_NOT_REAL_12345"]
    print("found")
except KeyError:
    print(True)

Solution

import os

home = os.environ.get("HOME") or os.environ.get("USERPROFILE")
if home:
    print("found_home")
else:
    print("no_home")

missing = os.environ.get("XYZZZY_NOT_REAL_12345", "default_used")
print(missing)

try:
    val = os.environ["XYZZZY_NOT_REAL_12345"]
    print("found")
except KeyError:
    print(True)

Output:

found_home
default_used
True

How it works: os.environ behaves like a dictionary that reflects the process environment. .get("KEY", default) returns the default if the key is missing, while os.environ["KEY"] raises KeyError. The HOME variable exists on macOS/Linux; USERPROFILE exists on Windows. The or chain finds whichever one exists.

Key insight: Always use os.environ.get() with a default value in production code rather than direct indexing. Direct indexing (os.environ["DB_URL"]) will crash your program if the variable is not set. The .get() approach lets you provide sensible defaults or explicitly handle the missing case. For required config, some teams prefer os.environ["KEY"] so the app fails fast with a clear error rather than silently using a wrong default.

Expected Output

found_home\ndefault_used\nTrue

Hints

Hint 1: os.environ.get() returns None (or a default) if the key does not exist — it never raises KeyError.

Hint 2: os.environ["KEY"] raises KeyError if the key is missing.

---

Medium

#5os.listdir vs os.scandir PerformanceMedium

os.listdiros.scandirDirEntryperformance

Demonstrate the difference between os.listdir() and os.scandir(). Show that scandir returns richer objects with cached file type information.

Python

import os
import tempfile

# Create a temp directory with some files and a subdirectory
with tempfile.TemporaryDirectory() as tmpdir:
    # Create test files
    open(os.path.join(tmpdir, "file1.txt"), "w").close()
    open(os.path.join(tmpdir, "file2.py"), "w").close()
    os.makedirs(os.path.join(tmpdir, "subdir"))

    # os.listdir returns plain strings
    names = os.listdir(tmpdir)
    print(all(isinstance(n, str) for n in names))

    # os.scandir returns DirEntry objects
    with os.scandir(tmpdir) as entries:
        entry_list = list(entries)

    print(all(hasattr(e, "is_file") for e in entry_list))

    # DirEntry can check file type without extra stat call
    files = [e.name for e in entry_list if e.is_file()]
    dirs = [e.name for e in entry_list if e.is_dir()]
    print(len(files) == 2 and len(dirs) == 1)

Solution

import os
import tempfile

with tempfile.TemporaryDirectory() as tmpdir:
    open(os.path.join(tmpdir, "file1.txt"), "w").close()
    open(os.path.join(tmpdir, "file2.py"), "w").close()
    os.makedirs(os.path.join(tmpdir, "subdir"))

    names = os.listdir(tmpdir)
    print(all(isinstance(n, str) for n in names))

    with os.scandir(tmpdir) as entries:
        entry_list = list(entries)

    print(all(hasattr(e, "is_file") for e in entry_list))

    files = [e.name for e in entry_list if e.is_file()]
    dirs = [e.name for e in entry_list if e.is_dir()]
    print(len(files) == 2 and len(dirs) == 1)

Output:

True
True
True

How it works:

1. os.listdir() returns a plain list of filename strings: ["file1.txt", "file2.py", "subdir"]. To check if each is a file or directory, you would need to call os.path.isfile() or os.path.isdir() — each of which makes a separate stat() system call.

2. os.scandir() returns DirEntry objects that cache the file type from the initial directory read. On Linux (getdents64) and Windows (FindNextFile), the OS already returns the file type in the directory listing — scandir exposes this without extra system calls.

3. The DirEntry objects provide .is_file(), .is_dir(), .is_symlink(), .stat(), .name, and .path — all without the per-file overhead of separate stat() calls.

Key insight: For directories with thousands of files, os.scandir() can be 2-20x faster than os.listdir() + os.path.isfile() because it avoids one stat() call per entry. Python 3.6+ uses scandir internally to implement os.walk(), which is why os.walk() also got dramatically faster in Python 3.6.

Expected Output

True\nTrue\nTrue

Hints

Hint 1: os.listdir() returns a list of name strings. os.scandir() returns DirEntry objects with cached metadata.

Hint 2: DirEntry.is_file() and .is_dir() do not require an extra stat() call on most operating systems.

#6os.walk: Recursive Directory TraversalMedium

os.walktraversaltopdown

Use os.walk() to traverse a directory tree and collect files at each level.

Python

import os
import tempfile

with tempfile.TemporaryDirectory() as root:
    # Build a tree:
    #   root/
    #     a.txt
    #     b.txt
    #     sub1/
    #       c.txt
    #     sub2/
    #       d.txt
    os.makedirs(os.path.join(root, "sub1"))
    os.makedirs(os.path.join(root, "sub2"))
    open(os.path.join(root, "a.txt"), "w").close()
    open(os.path.join(root, "b.txt"), "w").close()
    open(os.path.join(root, "sub1", "c.txt"), "w").close()
    open(os.path.join(root, "sub2", "d.txt"), "w").close()

    # Count total directories visited
    visited = list(os.walk(root))
    print(len(visited))

    # Print sorted filenames at each level
    for dirpath, dirnames, filenames in visited:
        print(sorted(filenames))

Solution

import os
import tempfile

with tempfile.TemporaryDirectory() as root:
    os.makedirs(os.path.join(root, "sub1"))
    os.makedirs(os.path.join(root, "sub2"))
    open(os.path.join(root, "a.txt"), "w").close()
    open(os.path.join(root, "b.txt"), "w").close()
    open(os.path.join(root, "sub1", "c.txt"), "w").close()
    open(os.path.join(root, "sub2", "d.txt"), "w").close()

    visited = list(os.walk(root))
    print(len(visited))

    for dirpath, dirnames, filenames in visited:
        print(sorted(filenames))

Output:

3
['a.txt', 'b.txt']
['c.txt']
['d.txt']

How it works: os.walk(root) is a generator that yields one tuple per directory in the tree. With the default topdown=True, it visits the root first, then descends into subdirectories alphabetically.

• First yield: (root, ["sub1", "sub2"], ["a.txt", "b.txt"]) — the root directory has 2 subdirs and 2 files.
• Second yield: (root/sub1, [], ["c.txt"]) — sub1 has no subdirs and 1 file.
• Third yield: (root/sub2, [], ["d.txt"]) — sub2 has no subdirs and 1 file.

Total: 3 directories visited.

Key insight: The topdown parameter controls traversal order. With topdown=True, you can modify the dirnames list in-place to prune the walk — for example, dirnames.remove("__pycache__") skips that subdirectory entirely. With topdown=False, children are visited before parents, which is useful when you need to delete directories bottom-up.

Expected Output

3\n['a.txt', 'b.txt']\n['c.txt']\n['d.txt']

Hints

Hint 1: os.walk() yields (dirpath, dirnames, filenames) for each directory in the tree.

Hint 2: With topdown=True (default), parent directories are visited before children.

#7os.makedirs with exist_okMedium

os.makedirsexist_okmkdir

Compare os.mkdir() vs os.makedirs() and demonstrate the exist_ok parameter.

Python

import os
import tempfile

with tempfile.TemporaryDirectory() as root:
    # makedirs creates the full chain of directories
    deep = os.path.join(root, "a", "b", "c")
    os.makedirs(deep)
    print(os.path.isdir(deep))

    # exist_ok=True allows calling makedirs on existing directory
    os.makedirs(deep, exist_ok=True)
    print(os.path.isdir(deep))

    # Without exist_ok, it raises FileExistsError
    try:
        os.makedirs(deep)
    except FileExistsError:
        print("FileExistsError")

    # os.mkdir only creates one level
    single = os.path.join(root, "single")
    os.mkdir(single)
    print(os.path.isdir(single))

Solution

import os
import tempfile

with tempfile.TemporaryDirectory() as root:
    deep = os.path.join(root, "a", "b", "c")
    os.makedirs(deep)
    print(os.path.isdir(deep))

    os.makedirs(deep, exist_ok=True)
    print(os.path.isdir(deep))

    try:
        os.makedirs(deep)
    except FileExistsError:
        print("FileExistsError")

    single = os.path.join(root, "single")
    os.mkdir(single)
    print(os.path.isdir(single))

Output:

True
True
FileExistsError
True

How it works:

1. os.makedirs(deep) creates the entire path a/b/c — including intermediate directories a and a/b. This is equivalent to mkdir -p in Unix.

2. Calling os.makedirs(deep, exist_ok=True) on an already-existing directory succeeds silently. This is the standard pattern for ensuring a directory exists before writing to it.

3. Without exist_ok=True, calling os.makedirs() on an existing directory raises FileExistsError.

4. os.mkdir(single) creates only a single directory level. If the parent does not exist, it raises FileNotFoundError.

Key insight: Always use os.makedirs(path, exist_ok=True) in production code — it handles both the "does not exist yet" and "already exists" cases in one call. The older pattern of checking if not os.path.exists(path): os.makedirs(path) has a race condition: another process could create the directory between the check and the creation.

Expected Output

True\nTrue\nFileExistsError\nTrue

Hints

Hint 1: os.makedirs() creates all intermediate directories. os.mkdir() creates only one level.

Hint 2: exist_ok=True suppresses the error when the directory already exists.

#8os.walk with Pruning: Skip DirectoriesMedium

os.walktopdownpruningdirnames

Use os.walk() with in-place pruning to skip directories whose names start with a dot or underscore.

Python

import os
import tempfile

with tempfile.TemporaryDirectory() as root:
    # Build tree with directories to skip
    os.makedirs(os.path.join(root, ".hidden"))
    os.makedirs(os.path.join(root, "__pycache__"))
    os.makedirs(os.path.join(root, "keep"))
    open(os.path.join(root, ".hidden", "secret.txt"), "w").close()
    open(os.path.join(root, "__pycache__", "cache.pyc"), "w").close()
    open(os.path.join(root, "keep", "good.txt"), "w").close()
    open(os.path.join(root, "keep.txt"), "w").close()

    visited = []
    for dirpath, dirnames, filenames in os.walk(root):
        # Prune: skip dirs starting with . or _
        dirnames[:] = [
            d for d in dirnames
            if not d.startswith(".") and not d.startswith("_")
        ]
        if filenames:
            visited.append(sorted(filenames))

    print(len(visited))
    for files in visited:
        print(files)

Solution

import os
import tempfile

with tempfile.TemporaryDirectory() as root:
    os.makedirs(os.path.join(root, ".hidden"))
    os.makedirs(os.path.join(root, "__pycache__"))
    os.makedirs(os.path.join(root, "keep"))
    open(os.path.join(root, ".hidden", "secret.txt"), "w").close()
    open(os.path.join(root, "__pycache__", "cache.pyc"), "w").close()
    open(os.path.join(root, "keep", "good.txt"), "w").close()
    open(os.path.join(root, "keep.txt"), "w").close()

    visited = []
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = [
            d for d in dirnames
            if not d.startswith(".") and not d.startswith("_")
        ]
        if filenames:
            visited.append(sorted(filenames))

    print(len(visited))
    for files in visited:
        print(files)

Output:

2
['keep.txt']
['good.txt']

How it works: The critical line is dirnames[:] = [...] — this modifies the list in-place. When topdown=True (the default), os.walk uses the same dirnames list to decide which subdirectories to descend into. By filtering out .hidden and __pycache__, we prevent os.walk from ever visiting those directories or their contents.

The walk visits only 2 directories with files: the root (containing keep.txt) and keep/ (containing good.txt). The .hidden/ and __pycache__/ directories and their files are completely skipped.

Key insight: You must use slice assignment (dirnames[:] = ...) or .remove() — a plain reassignment like dirnames = [...] creates a new local list and does not affect the one os.walk is tracking. This is one of the most common os.walk bugs. This pruning technique is essential for performance when traversing large directory trees — skipping .git/, node_modules/, and __pycache__/ can reduce traversal time by orders of magnitude.

Expected Output

2\n['keep.txt']\n['good.txt']

Hints

Hint 1: When topdown=True, modifying dirnames in-place controls which subdirectories os.walk descends into.

Hint 2: Use dirnames[:] = [...] or dirnames.remove() to prune — plain reassignment does not work.

---

Hard

#9Platform-Aware Path BuilderHard

os.nameos.sepplatform-detectioncross-platform

Build a cross-platform utility that detects the OS and constructs paths accordingly. Verify that os.path.join() always uses the correct separator.

Python

import os
import sys

# Detect platform
is_posix = os.name == "posix"
is_windows = os.name == "nt"
print(is_posix or is_windows)

# os.sep matches the platform
if is_posix:
    print(os.sep == "/")
else:
    print(os.sep == "\\")

# os.path.join uses the correct separator automatically
joined = os.path.join("home", "user", "documents", "file.txt")
print(os.sep in joined)

# os.pathsep separates PATH entries (: on POSIX, ; on Windows)
path_var = os.environ.get("PATH", "")
if is_posix:
    print(os.pathsep == ":")
else:
    print(os.pathsep == ";")

Solution

import os
import sys

is_posix = os.name == "posix"
is_windows = os.name == "nt"
print(is_posix or is_windows)

if is_posix:
    print(os.sep == "/")
else:
    print(os.sep == "\\")

joined = os.path.join("home", "user", "documents", "file.txt")
print(os.sep in joined)

path_var = os.environ.get("PATH", "")
if is_posix:
    print(os.pathsep == ":")
else:
    print(os.pathsep == ";")

Output:

True
True
True
True

How it works:

1. os.name returns "posix" on macOS and Linux, "nt" on Windows. Every mainstream OS is one of these two, so the or is always True.

2. os.sep is the directory separator for the current OS: / on POSIX, \\ on Windows. This matches the platform we detected.

3. os.path.join("home", "user", "documents", "file.txt") produces home/user/documents/file.txt on POSIX or home\user\documents\file.txt on Windows. The separator is always present in the result.

4. os.pathsep is the separator used in the PATH environment variable: : on POSIX (e.g., /usr/bin:/usr/local/bin), ; on Windows (e.g., C:\Windows;C:\Python).

Key insight: Never hardcode / or \\ as path separators. Use os.path.join() or pathlib.Path to construct paths, and os.sep/os.pathsep when you need to split or inspect path strings. Code that hardcodes separators breaks silently on the other platform — the files appear to exist as single-component filenames containing slashes.

Expected Output

True\nTrue\nTrue\nTrue

Hints

Hint 1: os.name returns "posix" on macOS/Linux and "nt" on Windows.

Hint 2: os.sep is "/" on POSIX systems and "\\" on Windows.

Hint 3: os.path.join() automatically uses the correct separator for the current OS.

#10Build a File Finder with os.walkHard

os.walkos.pathfile-searchrecursive

Implement a recursive file finder that searches a directory tree for files matching a given extension. Return full paths sorted alphabetically by filename.

Python

import os
import tempfile

def find_files_by_ext(root_dir, extension):
    """Find all files with a given extension recursively."""
    matches = []
    for dirpath, dirnames, filenames in os.walk(root_dir):
        # Skip hidden directories
        dirnames[:] = [d for d in dirnames if not d.startswith(".")]
        for fname in filenames:
            if fname.endswith(extension):
                full_path = os.path.join(dirpath, fname)
                matches.append(full_path)
    # Sort by filename only (not by full path)
    matches.sort(key=lambda p: os.path.basename(p))
    return matches

# Test it
with tempfile.TemporaryDirectory() as root:
    # Create a nested structure with mixed file types
    os.makedirs(os.path.join(root, "config"))
    os.makedirs(os.path.join(root, "src"))
    os.makedirs(os.path.join(root, ".git"))

    for name in ["settings.yaml", "README.md"]:
        open(os.path.join(root, name), "w").close()
    for name in ["db.yaml", "app.py"]:
        open(os.path.join(root, "config", name), "w").close()
    for name in ["config.yaml", "main.py"]:
        open(os.path.join(root, "src", name), "w").close()
    open(os.path.join(root, ".git", "config.yaml"), "w").close()

    results = find_files_by_ext(root, ".yaml")
    print(len(results))
    print([os.path.basename(r) for r in results])
    # Verify .git was skipped
    print(all(".git" not in r for r in results))

Solution

import os
import tempfile

def find_files_by_ext(root_dir, extension):
    matches = []
    for dirpath, dirnames, filenames in os.walk(root_dir):
        dirnames[:] = [d for d in dirnames if not d.startswith(".")]
        for fname in filenames:
            if fname.endswith(extension):
                full_path = os.path.join(dirpath, fname)
                matches.append(full_path)
    matches.sort(key=lambda p: os.path.basename(p))
    return matches

with tempfile.TemporaryDirectory() as root:
    os.makedirs(os.path.join(root, "config"))
    os.makedirs(os.path.join(root, "src"))
    os.makedirs(os.path.join(root, ".git"))

    for name in ["settings.yaml", "README.md"]:
        open(os.path.join(root, name), "w").close()
    for name in ["db.yaml", "app.py"]:
        open(os.path.join(root, "config", name), "w").close()
    for name in ["config.yaml", "main.py"]:
        open(os.path.join(root, "src", name), "w").close()
    open(os.path.join(root, ".git", "config.yaml"), "w").close()

    results = find_files_by_ext(root, ".yaml")
    print(len(results))
    print([os.path.basename(r) for r in results])
    print(all(".git" not in r for r in results))

Output:

3
['config.yaml', 'db.yaml', 'settings.yaml']
True

How it works: The function walks the entire directory tree, pruning hidden directories (.git, .venv, etc.) via dirnames[:] = [...]. For each directory visited, it checks every filename against the target extension and collects full paths. Results are sorted by basename for deterministic output.

The .git/config.yaml file is not found because .git was pruned from dirnames during the walk. The 3 matches are: src/config.yaml, config/db.yaml, and settings.yaml in the root.

Key insight: This pattern — walk + prune + filter + collect — is the foundation of most file search tools. In production, you would add more pruning rules (skip node_modules, __pycache__, .venv), support multiple extensions, and possibly use os.scandir for better performance. For simpler cases, pathlib.Path(root).rglob("*.yaml") achieves the same result in one line, though it does not support directory pruning.

Expected Output

3\n['config.yaml', 'db.yaml', 'settings.yaml']\nTrue

Hints

Hint 1: Use os.walk() to traverse all directories and filter filenames by extension.

Hint 2: os.path.splitext() gives you the extension to match against.

Hint 3: Sort results for deterministic output.

#11Environment Variable Safety: Isolation PatternHard

os.environisolationcontext-managerproduction

Build a context manager that temporarily sets environment variables and restores the original state on exit — even if an exception occurs.

Python

import os
from contextlib import contextmanager

@contextmanager
def temp_environ(**kwargs):
    """Temporarily set environment variables, restore on exit."""
    old_values = {}
    for key, value in kwargs.items():
        old_values[key] = os.environ.get(key)
        os.environ[key] = value
    try:
        yield
    finally:
        for key in kwargs:
            old_val = old_values[key]
            if old_val is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = old_val

# Test: modify existing variable
os.environ["TEST_VAR"] = "original"
print(os.environ["TEST_VAR"])

with temp_environ(TEST_VAR="temporary"):
    print(os.environ["TEST_VAR"])

print(os.environ["TEST_VAR"])

# Test: set a new variable that did not exist
with temp_environ(BRAND_NEW_VAR="hello"):
    pass

print(os.environ.get("BRAND_NEW_VAR"))

# Test: restoration works even after exception
try:
    with temp_environ(TEST_VAR="error_value"):
        raise ValueError("something broke")
except ValueError:
    pass

# Clean up
print(os.environ.get("BRAND_NEW_VAR"))
del os.environ["TEST_VAR"]

Solution

import os
from contextlib import contextmanager

@contextmanager
def temp_environ(**kwargs):
    old_values = {}
    for key, value in kwargs.items():
        old_values[key] = os.environ.get(key)
        os.environ[key] = value
    try:
        yield
    finally:
        for key in kwargs:
            old_val = old_values[key]
            if old_val is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = old_val

os.environ["TEST_VAR"] = "original"
print(os.environ["TEST_VAR"])

with temp_environ(TEST_VAR="temporary"):
    print(os.environ["TEST_VAR"])

print(os.environ["TEST_VAR"])

with temp_environ(BRAND_NEW_VAR="hello"):
    pass

print(os.environ.get("BRAND_NEW_VAR"))

try:
    with temp_environ(TEST_VAR="error_value"):
        raise ValueError("something broke")
except ValueError:
    pass

print(os.environ.get("BRAND_NEW_VAR"))
del os.environ["TEST_VAR"]

Output:

original
temporary
original
None
None

How it works:

1. Before entering the with block, temp_environ saves the current value of each key (or None if it does not exist), then sets the new values.

2. Inside the with block, os.environ["TEST_VAR"] reflects the temporary value "temporary".

3. When the with block exits (via the finally clause), the original values are restored. If the original value was None (the variable did not exist), it is removed from the environment entirely.

4. The finally block guarantees restoration even if an exception is raised inside the with block. After the ValueError, TEST_VAR is back to "original" and BRAND_NEW_VAR is removed.

Key insight: os.environ mutations are process-global and inherited by child processes. In testing and production, you must ensure temporary environment changes are always cleaned up. This context manager pattern is used by popular libraries like pytest (via monkeypatch.setenv) and unittest.mock.patch.dict(os.environ, ...). Without proper isolation, a test that sets DATABASE_URL could leak that value into subsequent tests, causing intermittent failures that are extremely hard to debug.

Expected Output

original\ntemporary\noriginal\nNone\nNone

Hints

Hint 1: os.environ modifications affect the entire process. Always restore original values after temporary changes.

Hint 2: A context manager pattern can automate save-and-restore for environment variables.