Export Verbosity as an IntEnum

Author: TylerYep

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2021 Tyler Yep
+Copyright (c) 2022 Tyler Yep
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

tests/torchinfo_test.py

@@ -29,6 +29,7 @@
 )
 from tests.fixtures.tmva_net import TMVANet  # type: ignore[attr-defined]
 from torchinfo import ColumnSettings, summary
+from torchinfo.enums import Verbosity
 
 
 def test_basic_summary() -> None:
@@ -253,7 +254,7 @@ def test_lstm() -> None:
         LSTMNet(),
         input_size=(1, 100),
         dtypes=[torch.long],
-        verbose=2,
+        verbose=Verbosity.VERBOSE,
         col_width=20,
         col_names=("kernel_size", "output_size", "num_params", "mult_adds"),
         row_settings=("var_names",),

torchinfo/__init__.py

@@ -1,7 +1,7 @@
 """ torchinfo """
-from .enums import ColumnSettings, RowSettings
+from .enums import ColumnSettings, RowSettings, Verbosity
 from .model_statistics import ModelStatistics
 from .torchinfo import summary
 
-__all__ = ("ModelStatistics", "summary", "ColumnSettings", "RowSettings")
-__version__ = "1.6.0"
+__all__ = ("summary", "ColumnSettings", "ModelStatistics", "RowSettings", "Verbosity")
+__version__ = "1.6.1"

torchinfo/enums.py

@@ -1,7 +1,7 @@
 """ constants.py """
 from __future__ import annotations
 
-from enum import Enum, unique
+from enum import Enum, IntEnum, unique
 
 
 @unique
@@ -25,7 +25,7 @@ class ColumnSettings(Enum):
 
 
 @unique
-class Verbosity(Enum):
+class Verbosity(IntEnum):
     """Contains verbosity levels."""
 
     QUIET, DEFAULT, VERBOSE = 0, 1, 2

torchinfo/formatting.py

@@ -119,7 +119,7 @@ def layer_info_to_row(
         layer_name = layer_info.get_layer_name(self.show_var_name, self.show_depth)
         new_line = self.format_row(f"{start_str}{layer_name}", values_for_row)
 
-        if self.verbose == Verbosity.VERBOSE.value:
+        if self.verbose == Verbosity.VERBOSE:
             for inner_name, inner_layer_info in layer_info.inner_layers.items():
                 prefix = self.get_start_str(layer_info.depth + 1)
                 new_line += self.format_row(f"{prefix}{inner_name}", inner_layer_info)

torchinfo/torchinfo.py

@@ -205,7 +205,7 @@ class name as the key. If the forward pass is an expensive operation,
     results = ModelStatistics(
         summary_list, correct_input_size, get_total_memory_used(x), formatting
     )
-    if verbose > Verbosity.QUIET.value:
+    if verbose > Verbosity.QUIET:
         print(results)
     return results