[Python 3.10+] Modernize codebase with Python 3.10+ idioms and features

[rahul-tuli]

Now that we've dropped Python 3.9 support (#486 ), we can modernize our codebase to use Python 3.10+ features for better readability and maintainability. This is a great opportunity for community contributions!

Background

With Python 3.10+, we have access to several powerful features that make our code more readable and maintainable:

• PEP 604: Union types can use X | Y instead of Union[X, Y]
• PEP 585: Built-in collection types (list, dict, tuple) can be used directly instead of importing from typing
• PEP 634: Structural pattern matching with match/case statements

Currently, compressed-tensors has:

• 132 instances of Optional[] across the codebase
• 47 instances of Union[] across the codebase
• 179 total type hint instances that could be modernized across 34 files
• 31 isinstance checks that could potentially benefit from pattern matching

This is a great opportunity for community contributions! This issue tracks the effort to modernize our type hints and leverage Python 3.10+ features.

---

Modernization Categories

1. Type Hints with | Operator

Replace Union[] and Optional[] with the cleaner | syntax.

Example from src/compressed_tensors/quantization/lifecycle/forward.py:50-53:

# Before (Python 3.9 style)
from typing import Optional

def quantize(
    x: torch.Tensor,
    scale: torch.Tensor,
    zero_point: torch.Tensor,
    args: QuantizationArgs,
    dtype: Optional[torch.dtype] = None,
    g_idx: Optional[torch.Tensor] = None,
    global_scale: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    ...

# After (Python 3.10+ style)
def quantize(
    x: torch.Tensor,
    scale: torch.Tensor,
    zero_point: torch.Tensor,
    args: QuantizationArgs,
    dtype: torch.dtype | None = None,
    g_idx: torch.Tensor | None = None,
    global_scale: torch.Tensor | None = None,
) -> torch.Tensor:
    ...

Example from src/compressed_tensors/utils/offload.py:150:

# Before
def cast_to_device(device_spec: Union[int, torch.device]) -> torch.device:
    ...

# After
def cast_to_device(device_spec: int | torch.device) -> torch.device:
    ...

2. Built-in Generic Types

Use built-in types like list, dict, tuple instead of importing from typing.

Example from src/compressed_tensors/utils/match.py:18:

# Before
from typing import List, Tuple, Union, Iterable

def match_named_modules(
    model: torch.nn.Module,
    targets: Optional[Iterable[str]],
    ignore: Optional[Iterable[str]] = None,
) -> Generator[Tuple[str, torch.nn.Module]]:
    ...

# After  
from collections.abc import Generator, Iterable

def match_named_modules(
    model: torch.nn.Module,
    targets: Iterable[str] | None,
    ignore: Iterable[str] | None = None,
) -> Generator[tuple[str, torch.nn.Module]]:
    ...

Example from src/compressed_tensors/quantization/quant_args.py:174:

# Before
from typing import Union, List

class QuantizationArgs(BaseModel):
    block_structure: Optional[List[int]] = None
    dynamic: Union[DynamicType, bool] = False
    ...

# After
class QuantizationArgs(BaseModel):
    block_structure: list[int] | None = None
    dynamic: DynamicType | bool = False
    ...

3. Structural Pattern Matching

Replace isinstance chains with match/case statements for clearer intent.

Example from src/compressed_tensors/utils/offload.py:305-342:

# Before
def offload_to_weights_map(
    weights_map: Union[PrefixedDataset, Dict, OffloadedWeightsLoader],
    key: str,
    value: torch.Tensor,
    offload_device: Optional[Union[torch.device, Literal["disk"]]] = None,
):
    if isinstance(weights_map, PrefixedDataset):
        if offload_device == "disk":
            raise ValueError(f"Cannot offload to disk with type {type(weights_map)}")
        dataset = weights_map.dataset
        key = f"{weights_map.prefix}{key}"
        offload_to_weights_map(dataset, key, value, offload_device)
        
    elif isinstance(weights_map, OffloadedWeightsLoader):
        if key not in weights_map.all_keys:
            weights_map.all_keys.append(key)
        if len(weights_map.index) <= 0 and offload_device != "disk":
            offload_to_weights_map(weights_map.state_dict, key, value, offload_device)
        else:
            raise NotImplementedError(
                "Updating weights_map with disk offloading is not implemented yet"
            )
            
    elif isinstance(weights_map, dict):
        if offload_device == "disk":
            raise ValueError(f"Cannot offload to disk with type {type(weights_map)}")
        # ... rest of logic
    else:
        raise NotImplementedError(
            "Updating offload data not implemented for weights_map of type "
            f"{type(weights_map)}"
        )

# After
def offload_to_weights_map(
    weights_map: PrefixedDataset | dict | OffloadedWeightsLoader,
    key: str,
    value: torch.Tensor,
    offload_device: torch.device | Literal["disk"] | None = None,
):
    match weights_map:
        case PrefixedDataset():
            if offload_device == "disk":
                raise ValueError(f"Cannot offload to disk with type {type(weights_map)}")
            dataset = weights_map.dataset
            key = f"{weights_map.prefix}{key}"
            offload_to_weights_map(dataset, key, value, offload_device)
            
        case OffloadedWeightsLoader():
            if key not in weights_map.all_keys:
                weights_map.all_keys.append(key)
            if len(weights_map.index) <= 0 and offload_device != "disk":
                offload_to_weights_map(weights_map.state_dict, key, value, offload_device)
            else:
                raise NotImplementedError(
                    "Updating weights_map with disk offloading is not implemented yet"
                )
                
        case dict():
            if offload_device == "disk":
                raise ValueError(f"Cannot offload to disk with type {type(weights_map)}")
            # ... rest of logic
            
        case _:
            raise NotImplementedError(
                "Updating offload data not implemented for weights_map of type "
                f"{type(weights_map)}"
            )

---

How to Contribute

Getting Started

1. Fork and clone the repository
2. Create a new branch for your changes: git checkout -b modernize-python310-<module-name>
3. Make your changes following the examples above
4. Run tests and quality checks (see below)
5. Submit a PR referencing this issue

Contribution Sizes

Small PRs (Recommended for first-time contributors):

• Modernize 1-2 files with 10-20 type hint changes
• Focus on straightforward Union[] → | and Optional[] → | None conversions
• Examples: Helper files, utility modules

Medium PRs:

• Modernize a complete module (e.g., all files in src/compressed_tensors/quantization/)
• Include both type hints and pattern matching improvements
• 3-5 files with mixed changes

Large PRs:

• Complete modernization of a major subsystem
• Comprehensive pattern matching refactoring
• Multiple modules with coordinated changes

Suggested Files to Start With

Easy (Good first issues):

• src/compressed_tensors/quantization/lifecycle/forward.py (17 Optional instances)
• src/compressed_tensors/utils/match.py (11 Optional instances)
• src/compressed_tensors/quantization/utils/helpers.py (7 Optional instances)

Medium:

• src/compressed_tensors/utils/offload.py (8 Union + 5 Optional instances, plus isinstance chains)
• src/compressed_tensors/compressors/model_compressors/model_compressor.py (12 Optional + 8 Union instances)
• src/compressed_tensors/quantization/quant_args.py (7 Optional + 5 Union instances)

Advanced:

• src/compressed_tensors/config/format.py (Complex type hierarchies)
• Pattern matching refactoring in compressor modules

---

Requirements for PRs

✅ Must have:

• All make quality checks pass (ruff formatting and linting)
• Relevant tests pass (pytest tests/{module} -v)
• No functional changes (type hints/style only)
• Clean commit messages
• Reference this issue (e.g., "Part of #XXX")

✅ Nice to have:

• Updated docstrings if they reference old type syntax
• Multiple files in same module (for consistency)
• Comments explaining complex pattern matches

---

Testing Guidelines

Before submitting your PR, ensure all checks pass:

# Run code quality checks
make quality

# Run tests
make test

# Or run both
make quality && make test

Module-specific testing:

# Test specific modules
pytest tests/test_quantization/ -v
pytest tests/test_compression/ -v

---

Progress Tracking

Overall Progress

• Type hints modernization (Union[] and Optional[] → |)
• Built-in generic types (List, Dict, Tuple → list, dict, tuple)
• Pattern matching (isinstance chains → match/case)

By Module

Quantization Module:

• src/compressed_tensors/quantization/lifecycle/forward.py
• src/compressed_tensors/quantization/quant_args.py
• src/compressed_tensors/quantization/utils/helpers.py
• src/compressed_tensors/quantization/quant_config.py

Compressors Module:

• src/compressed_tensors/compressors/model_compressors/model_compressor.py
• src/compressed_tensors/compressors/quantized_compressors/pack_quantized.py
• src/compressed_tensors/compressors/quantized_compressors/nvfp4_quantized.py
• src/compressed_tensors/compressors/quantized_compressors/naive_quantized.py
• src/compressed_tensors/compressors/sparse_compressors/sparse_24_bitmask.py

Utilities Module:

• src/compressed_tensors/utils/offload.py
• src/compressed_tensors/utils/match.py
• src/compressed_tensors/utils/safetensors_load.py

Config Module:

• src/compressed_tensors/config/format.py

Registry Module:

• src/compressed_tensors/registry/registry.py

---

Resources

Python Enhancement Proposals (PEPs)

• PEP 604 - Union Type Operator
• PEP 585 - Type Hinting Generics In Standard Collections
• PEP 634 - Structural Pattern Matching

Guides and Tutorials

• Real Python - Python 3.10 New Features
• Python Match Statement Guide
• Type Hints Cheat Sheet

---

Example PRs

As PRs are merged, we'll link them here as examples:

• (To be added as PRs are created)

---

Questions?

Feel free to ask questions in this issue or in your PR! Common questions:

Q: Should I convert all type hints in a file at once?
A: Yes, for consistency. If you're modernizing a file, update all the type hints in that file.

Q: What about forward references and circular imports?
A: You can still use string annotations "ClassName" or use from __future__ import annotations at the top of the file.

Q: Should I use pattern matching for all isinstance checks?
A: Not necessarily. Use pattern matching when it improves readability, especially for chains of 3+ isinstance checks or when matching on structure/attributes.

Q: Can I combine multiple small files into one PR?
A: Yes! Just make sure they're related (e.g., all utility files or all in the same module).

---

Let's modernize compressed-tensors together! 🚀