Move minimum version to Python 3.7

Author: TylerYep

.pre-commit-config.yaml

@@ -27,6 +27,12 @@ repos:
     rev: 4.0.1
     hooks:
       - id: flake8
+        additional_dependencies:
+          [
+            flake8-future-annotations,
+            flake8-bugbear,
+            flake8-comprehensions,
+          ]
 
   - repo: local
     hooks:

README.md

@@ -1,6 +1,6 @@
 # torchinfo
 
-[![Python 3.6+](https://img.shields.io/badge/python-3.6+-blue.svg)](https://www.python.org/downloads/release/python-360/)
+[![Python 3.7+](https://img.shields.io/badge/python-3.7+-blue.svg)](https://www.python.org/downloads/release/python-360/)
 [![PyPI version](https://badge.fury.io/py/torchinfo.svg)](https://badge.fury.io/py/torchinfo)
 [![Conda version](https://img.shields.io/conda/vn/conda-forge/torchinfo)](https://anaconda.org/conda-forge/torchinfo)
 [![Build Status](https://github.com/TylerYep/torchinfo/actions/workflows/test.yml/badge.svg)](https://github.com/TylerYep/torchinfo/actions/workflows/test.yml)
@@ -22,6 +22,7 @@ pip install torchinfo
 ```
 
 Alternatively, via conda:
+
 ```
 conda install -c conda-forge torchinfo
 ```
@@ -452,7 +453,7 @@ Estimated Total Size (MB): 0.00
 All issues and pull requests are much appreciated! If you are wondering how to build the project:
 
 - torchinfo is actively developed using the lastest version of Python.
-  - Changes should be backward compatible with Python 3.6, but this is subject to change in the future.
+  - Changes should be backward compatible to Python 3.7, and will follow Python's End-of-Life guidance for old versions.
   - Run `pip install -r requirements-dev.txt`. We use the latest versions of all dev packages.
   - Run `pre-commit install`.
   - To use auto-formatting tools, use `pre-commit run -a`.

setup.cfg

@@ -23,7 +23,7 @@ keywords = torch pytorch torchsummary torch-summary summary keras deep-learning
 
 [options]
 packages = torchinfo
-python_requires = >=3.6
+python_requires = >=3.7
 include_package_data = True
 
 [options.package_data]

tests/fixtures/models.py

@@ -1,7 +1,9 @@
 """ fixtures/models.py """
+from __future__ import annotations
+
 import math
 from collections import namedtuple
-from typing import Any, Dict, Tuple, cast
+from typing import Any, cast
 
 import torch
 from torch import nn
@@ -113,7 +115,7 @@ def __init__(
         self.encoder = nn.LSTM(embed_dim, hidden_dim, num_layers=num_layers)  # type: ignore[no-untyped-call] # noqa
         self.decoder = nn.Linear(hidden_dim, vocab_size)
 
-    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         embed = self.embedding(x)
         out, hidden = self.encoder(embed)
         out = self.decoder(out)
@@ -244,7 +246,7 @@ def __init__(self) -> None:
         self.fc1 = nn.Linear(320, 50)
         self.fc2 = nn.Linear(50, 10)
 
-    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
+    def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
         activation_dict = {}
         x = self.conv1(x)
         activation_dict["conv1"] = x
@@ -269,9 +271,9 @@ def __init__(self) -> None:
         super().__init__()
         self.return_dict = ReturnDictLayer()
 
-    def forward(self, x: torch.Tensor, y: Any) -> Dict[str, torch.Tensor]:
+    def forward(self, x: torch.Tensor, y: Any) -> dict[str, torch.Tensor]:
         del y
-        activation_dict: Dict[str, torch.Tensor] = self.return_dict(x)
+        activation_dict: dict[str, torch.Tensor] = self.return_dict(x)
         return activation_dict
 
 
@@ -282,7 +284,7 @@ def __init__(self) -> None:
         super().__init__()
         self.constant = 5
 
-    def forward(self, x: Dict[int, torch.Tensor], scale_factor: int) -> torch.Tensor:
+    def forward(self, x: dict[int, torch.Tensor], scale_factor: int) -> torch.Tensor:
         return scale_factor * (x[256] + x[512][0]) * self.constant
 
 
@@ -427,7 +429,7 @@ def __init__(self) -> None:
         self.parameter = torch.rand(3, 3, requires_grad=True)
         self.example_input_array = torch.zeros(1, 2, 3, 4, 5)
 
-    def forward(self) -> Dict[str, Any]:
+    def forward(self) -> dict[str, Any]:
         return {"loss": self.parameter.sum()}
 
 

torchinfo/formatting.py

@@ -1,7 +1,9 @@
 """ formatting.py """
+from __future__ import annotations
+
 import math
 from enum import Enum, unique
-from typing import Any, Dict, Iterable, List
+from typing import Any, Iterable
 
 from .layer_info import LayerInfo
 
@@ -64,8 +66,8 @@ def str_(val: Any) -> str:
 
     @staticmethod
     def get_children_layers(
-        summary_list: List[LayerInfo], layer_info: LayerInfo, index: int
-    ) -> List[LayerInfo]:
+        summary_list: list[LayerInfo], layer_info: LayerInfo, index: int
+    ) -> list[LayerInfo]:
         """Fetches all of the children of a given layer."""
         num_children = 0
         for layer in summary_list[index + 1 :]:
@@ -75,7 +77,7 @@ def get_children_layers(
         return summary_list[index + 1 : index + 1 + num_children]
 
     def set_layer_name_width(
-        self, summary_list: List[LayerInfo], align_val: int = 5
+        self, summary_list: list[LayerInfo], align_val: int = 5
     ) -> None:
         """
         Set layer name width by taking the longest line length and rounding up to
@@ -93,7 +95,7 @@ def get_total_width(self) -> int:
         """Calculate the total width of all lines in the table."""
         return len(tuple(self.col_names)) * self.col_width + self.layer_name_width
 
-    def format_row(self, layer_name: str, row_values: Dict[str, str]) -> str:
+    def format_row(self, layer_name: str, row_values: dict[str, str]) -> str:
         """Get the string representation of a single layer of the model."""
         info_to_use = [row_values.get(row_type, "") for row_type in self.col_names]
         new_line = f"{layer_name:<{self.layer_name_width}} "
@@ -113,7 +115,7 @@ def layer_info_to_row(
         self,
         layer_info: LayerInfo,
         reached_max_depth: bool,
-        children_layers: List["LayerInfo"],
+        children_layers: list[LayerInfo],
     ) -> str:
         """Convert layer_info to string representation of a row."""
         row_values = {
@@ -133,13 +135,13 @@ def layer_info_to_row(
                 new_line += self.format_row(f"{prefix}{inner_name}", inner_layer_info)
         return new_line
 
-    def layers_to_str(self, summary_list: List[LayerInfo]) -> str:
+    def layers_to_str(self, summary_list: list[LayerInfo]) -> str:
         """
         Print each layer of the model using a fancy branching diagram.
         This is necessary to handle Container modules that don't have explicit parents.
         """
         new_str = ""
-        current_hierarchy: Dict[int, LayerInfo] = {}
+        current_hierarchy: dict[int, LayerInfo] = {}
         for i, layer_info in enumerate(summary_list):
             if layer_info.depth > self.max_depth:
                 continue

torchinfo/layer_info.py

@@ -1,5 +1,7 @@
 """ layer_info.py """
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Union
+from __future__ import annotations
+
+from typing import Any, Dict, Iterable, Sequence, Union
 
 import torch
 from torch import nn
@@ -18,8 +20,8 @@ def __init__(
         var_name: str,
         module: nn.Module,
         depth: int,
-        depth_index: Optional[int] = None,
-        parent_info: Optional["LayerInfo"] = None,
+        depth_index: int | None = None,
+        parent_info: LayerInfo | None = None,
     ) -> None:
         # Identifying information
         self.layer_id = id(module)
@@ -30,7 +32,7 @@ def __init__(
             else module.__class__.__name__
         )
         # {layer name: {row_name: row_value}}
-        self.inner_layers: Dict[str, Dict[str, Any]] = {}
+        self.inner_layers: dict[str, dict[str, Any]] = {}
         self.depth = depth
         self.depth_index = depth_index
         self.executed = False
@@ -41,9 +43,9 @@ def __init__(
         # Statistics
         self.trainable_params = 0
         self.is_recursive = False
-        self.input_size: List[int] = []
-        self.output_size: List[int] = []
-        self.kernel_size: List[int] = []
+        self.input_size: list[int] = []
+        self.output_size: list[int] = []
+        self.kernel_size: list[int] = []
         self.num_params = 0
         self.macs = 0
 
@@ -52,11 +54,11 @@ def __repr__(self) -> str:
 
     @staticmethod
     def calculate_size(
-        inputs: DETECTED_INPUT_OUTPUT_TYPES, batch_dim: Optional[int]
-    ) -> List[int]:
+        inputs: DETECTED_INPUT_OUTPUT_TYPES, batch_dim: int | None
+    ) -> list[int]:
         """Set input_size or output_size using the model's inputs."""
 
-        def nested_list_size(inputs: Sequence[Any]) -> List[int]:
+        def nested_list_size(inputs: Sequence[Any]) -> list[int]:
             """Flattens nested list size."""
             if hasattr(inputs, "tensors"):
                 return nested_list_size(inputs.tensors)  # type: ignore[attr-defined]
@@ -164,7 +166,7 @@ def calculate_macs(self) -> None:
             elif "weight" in name or "bias" in name:
                 self.macs += prod(self.output_size[:2]) * param.nelement()
 
-    def check_recursive(self, summary_list: List["LayerInfo"]) -> None:
+    def check_recursive(self, summary_list: list[LayerInfo]) -> None:
         """
         If the current module is already-used, mark as (recursive).
         Must check before adding line to the summary.
@@ -175,7 +177,7 @@ def check_recursive(self, summary_list: List["LayerInfo"]) -> None:
                     self.is_recursive = True
 
     def macs_to_str(
-        self, reached_max_depth: bool, children_layers: List["LayerInfo"]
+        self, reached_max_depth: bool, children_layers: list[LayerInfo]
     ) -> str:
         """Convert MACs to string."""
         if self.macs <= 0:
@@ -199,7 +201,7 @@ def num_params_to_str(self, reached_max_depth: bool) -> str:
         return "--"
 
 
-def prod(num_list: Union[Iterable[int], torch.Size]) -> int:
+def prod(num_list: Iterable[int] | torch.Size) -> int:
     result = 1
     if isinstance(num_list, Iterable):
         for item in num_list:

torchinfo/model_statistics.py

@@ -1,5 +1,7 @@
 """ model_statistics.py """
-from typing import Any, List, Tuple
+from __future__ import annotations
+
+from typing import Any
 
 from .formatting import FormattingOptions
 from .layer_info import LayerInfo, prod
@@ -10,7 +12,7 @@ class ModelStatistics:
 
     def __init__(
         self,
-        summary_list: List[LayerInfo],
+        summary_list: list[LayerInfo],
         input_size: Any,
         total_input_size: int,
         formatting: FormattingOptions,
@@ -78,7 +80,7 @@ def to_megabytes(num: int) -> float:
         return num / 1e6
 
     @staticmethod
-    def to_readable(num: int) -> Tuple[str, float]:
+    def to_readable(num: int) -> tuple[str, float]:
         """Converts a number to millions, billions, or trillions."""
         if num >= 1e12:
             return "T", num / 1e12