[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! 🚀

[Agentx402]

Great issue! This is a well-structured modernization effort that will significantly improve code readability. I'd recommend starting with the "Easy" files mentioned (like forward.py or match.py) to get familiar with the patterns, and running make quality frequently as you work to catch any formatting issues early. The detailed examples and clear contribution guidelines make this very approachable for contributors at any level! 🚀

[ArkaSanka]

Hi @rahul-tuli, I am interested in working on this, please assign this to me.

[rahul-tuli]

Hi @rahul-tuli, I am interested in working on this, please assign this to me.

Hi @ArkaSanka Thanks for your interest, since this task will touch quite a few files, we'd like to start small, could you pick a module (file or folder) that you'd like to modernize first. Then comment here that you'll be working on modernizing that module (for book keeping) when the PR is ready please link this issue in the PR description. For assignments we'll use FCFS; seems like you have the first mover advantage on that

[ArkaSanka]

Sure, I'll be working on the Quantization module

[ArkaSanka]

In src/compressed_tensors/quantization/utils/helpers.py, there are a few Generators defined using the datatype from typing import Generator. Is it okay to change it to collections.abc.Generator to be consistent with PEP 585?
Can follow the same in other instances as well

[HDCharles]

Rahul is out today but yeah that file is the last one with from typing import generator so it would be good to fix