Renamed default_storage_type to default_storage_field.

ami-iit · Mar 13, 2023 · 67c68b9 · 67c68b9
1 parent 09ecc8e
commit 67c68b9
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 9 deletions.
diff --git a/src/hippopt/__init__.py b/src/hippopt/__init__.py
@@ -3,7 +3,7 @@
     OptimizationObject,
     StorageType,
     TOptimizationObject,
-    default_storage_type,
+    default_storage_field,
 )
 from .base.parameter import Parameter, TParameter
 from .base.variable import TVariable, Variable
diff --git a/src/hippopt/base/optimization_object.py b/src/hippopt/base/optimization_object.py
@@ -52,8 +52,8 @@ def get_default_initialized_object(
         return output
 
 
-def default_storage_type(input_type: Type[OptimizationObject]):
+def default_storage_field(cls: Type[OptimizationObject]):
     return dataclasses.field(
         default=None,
-        metadata=input_type.StorageTypeMetadata,
+        metadata=cls.StorageTypeMetadata,
     )
diff --git a/test/test_base.py b/test/test_base.py
@@ -8,21 +8,21 @@
     StorageType,
     TOptimizationObject,
     Variable,
-    default_storage_type,
+    default_storage_field,
 )
 
 
 @dataclasses.dataclass
 class TestVariable(OptimizationObject):
-    storage: StorageType = default_storage_type(Variable)
+    storage: StorageType = default_storage_field(cls=Variable)
 
     def __post_init__(self):
         self.storage = np.ones(shape=3)
 
 
 @dataclasses.dataclass
 class TestParameter(OptimizationObject):
-    storage: StorageType = default_storage_type(Parameter)
+    storage: StorageType = default_storage_field(cls=Parameter)
 
     def __post_init__(self):
         self.storage = np.ones(shape=3)
@@ -44,8 +44,8 @@ def test_zero_parameter():
 
 @dataclasses.dataclass
 class CustomInitializationVariable(OptimizationObject):
-    variable: StorageType = default_storage_type(Variable)
-    parameter: StorageType = default_storage_type(Parameter)
+    variable: StorageType = default_storage_field(cls=Variable)
+    parameter: StorageType = default_storage_field(cls=Parameter)
 
     def __post_init__(self):
         self.variable = np.ones(shape=3)
@@ -72,7 +72,7 @@ def test_custom_initialization():
 @dataclasses.dataclass
 class AggregateClass(OptimizationObject):
     aggregated: CustomInitializationVariable
-    other_parameter: StorageType = default_storage_type(Parameter)
+    other_parameter: StorageType = default_storage_field(cls=Parameter)
     other: str = ""
 
     def __post_init__(self):