Merge pull request #101 from VukManojlovic/CTX-5033

CTX-5033: Allow integers instead of MetricType enum as type arguments for Metric.create
coretex-ai · Dec 5, 2023 · 55a6e9b · 55a6e9b
2 parents 3370a1b + 289c897
commit 55a6e9b
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/coretex/entities/task_run/metrics/metric.py b/coretex/entities/task_run/metrics/metric.py
@@ -15,7 +15,7 @@
 #     You should have received a copy of the GNU Affero General Public License
 #     along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 from typing_extensions import Self
 
 from .metric_type import MetricType
@@ -44,9 +44,9 @@ def create(
         cls,
         name: str,
         xLabel: str,
-        xType: MetricType,
+        xType: Union[MetricType, int],
         yLabel: str,
-        yType: MetricType,
+        yType: Union[MetricType, int],
         xRange: Optional[List[float]] = None,
         yRange: Optional[List[float]] = None
     ) -> Self:
@@ -56,18 +56,24 @@ def create(
                 name of Metric
             xLabel : str
                 label of x axis which will be displayed
-            xType : MetricType
+            xType : Union[MetricType, int]
                 type of x axis which will be displayed
             yLabel : str
                 label of y axis which will be displayed
-            yType : MetricType
+            yType : Union[MetricType, int]
                 type of y axis which will be displayed
             xRange : Optional[List[float]]
                 range in which values will be displayed for x axis
             yRange : Optional[List[float]]
                 range in which values will be displayed for y axis
         """
 
+        if isinstance(xType, int):
+            xType = MetricType(xType)
+
+        if isinstance(yType, int):
+            yType = MetricType(yType)
+
         obj = cls()
 
         obj.name = name