Make use of Literal in back-end schema for vfp type (#875)

equinor · Feb 12, 2025 · 93303d5 · 93303d5
1 parent afccefb
commit 93303d5
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 39 deletions.
diff --git a/backend_py/primary/primary/routers/vfp/schemas.py b/backend_py/primary/primary/routers/vfp/schemas.py
@@ -1,43 +1,44 @@
+from enum import StrEnum
+from typing import Literal
+
 from pydantic import BaseModel
 
 from primary.services.sumo_access.vfp_types import THP, WFR, GFR, ALQ, FlowRateType, UnitType, TabType
 
 
-class VfpProdTable(BaseModel):
-    isProdTable: bool = True
+class VfpType(StrEnum):
+    PROD = "PROD"
+    INJ = "INJ"
+
+
+class VfpTableBase(BaseModel):
+    vfpType: Literal[VfpType.INJ, VfpType.PROD]
     tableNumber: int
     datum: float
-    thpType: THP
-    wfrType: WFR
-    gfrType: GFR
-    alqType: ALQ
     flowRateType: FlowRateType
     unitType: UnitType
     tabType: TabType
     thpValues: list[float]
-    wfrValues: list[float]
-    gfrValues: list[float]
-    alqValues: list[float]
     flowRateValues: list[float]
     bhpValues: list[float]
     flowRateUnit: str
     thpUnit: str
+    bhpUnit: str
+
+
+class VfpProdTable(VfpTableBase):
+    vfpType: Literal[VfpType.PROD] = VfpType.PROD
+    thpType: THP
+    wfrType: WFR
+    gfrType: GFR
+    alqType: ALQ
+    wfrValues: list[float]
+    gfrValues: list[float]
+    alqValues: list[float]
     wfrUnit: str
     gfrUnit: str
     alqUnit: str
-    bhpUnit: str
 
 
-class VfpInjTable(BaseModel):
-    isInjTable: bool = True
-    tableNumber: int
-    datum: float
-    flowRateType: FlowRateType
-    unitType: UnitType
-    tabType: TabType
-    thpValues: list[float]
-    flowRateValues: list[float]
-    bhpValues: list[float]
-    flowRateUnit: str
-    thpUnit: str
-    bhpUnit: str
+class VfpInjTable(VfpTableBase):
+    vfpType: Literal[VfpType.INJ] = VfpType.INJ
diff --git a/frontend/src/api/autogen/types.gen.ts b/frontend/src/api/autogen/types.gen.ts
@@ -890,7 +890,7 @@ export type VectorStatisticSensitivityData_api = {
 };
 
 export type VfpInjTable_api = {
-    isInjTable: boolean;
+    vfpType: "INJ";
     tableNumber: number;
     datum: number;
     flowRateType: FlowRateType_api;
@@ -905,28 +905,28 @@ export type VfpInjTable_api = {
 };
 
 export type VfpProdTable_api = {
-    isProdTable: boolean;
+    vfpType: "PROD";
     tableNumber: number;
     datum: number;
-    thpType: Thp_api;
-    wfrType: Wfr_api;
-    gfrType: Gfr_api;
-    alqType: Alq_api;
     flowRateType: FlowRateType_api;
     unitType: UnitType_api;
     tabType: TabType_api;
     thpValues: Array<number>;
-    wfrValues: Array<number>;
-    gfrValues: Array<number>;
-    alqValues: Array<number>;
     flowRateValues: Array<number>;
     bhpValues: Array<number>;
     flowRateUnit: string;
     thpUnit: string;
+    bhpUnit: string;
+    thpType: Thp_api;
+    wfrType: Wfr_api;
+    gfrType: Gfr_api;
+    alqType: Alq_api;
+    wfrValues: Array<number>;
+    gfrValues: Array<number>;
+    alqValues: Array<number>;
     wfrUnit: string;
     gfrUnit: string;
     alqUnit: string;
-    bhpUnit: string;
 };
 
 export enum Wfr_api {

diff --git a/frontend/src/modules/Vfp/utils/vfpDataAccessor.ts b/frontend/src/modules/Vfp/utils/vfpDataAccessor.ts
@@ -1,6 +1,6 @@
 import { Alq_api, FlowRateType_api, Gfr_api, VfpInjTable_api, VfpProdTable_api, Wfr_api } from "@api";
 
-import { isProdTable } from "./vfpTableClassifier";
+import { isInjTable, isProdTable } from "./vfpTableClassifier";
 
 import { VfpParam, VfpType } from "../types";
 
@@ -16,7 +16,7 @@ export class VfpDataAccessor {
     }
 
     isInjTable(): boolean {
-        return !isProdTable(this._vfpTable);
+        return isInjTable(this._vfpTable);
     }
 
     getTableNumber(): number {
@@ -65,7 +65,7 @@ export class VfpDataAccessor {
         if (vfpParam === VfpParam.THP) {
             return this._vfpTable.thpValues;
         }
-        if ("isProdTable" in this._vfpTable) {
+        if (isProdTable(this._vfpTable)) {
             if (vfpParam === VfpParam.WFR) {
                 return this._vfpTable.wfrValues;
             }

diff --git a/frontend/src/modules/Vfp/utils/vfpTableClassifier.ts b/frontend/src/modules/Vfp/utils/vfpTableClassifier.ts
@@ -1,5 +1,9 @@
-import { VfpProdTable_api, VfpInjTable_api } from "@api";
+import { VfpInjTable_api, VfpProdTable_api } from "@api";
 
 export function isProdTable(vfpTable: VfpProdTable_api | VfpInjTable_api): vfpTable is VfpProdTable_api {
-    return "isProdTable" in vfpTable
-}
+    return vfpTable.vfpType === "PROD";
+}
+
+export function isInjTable(vfpTable: VfpProdTable_api | VfpInjTable_api): vfpTable is VfpInjTable_api {
+    return vfpTable.vfpType === "INJ";
+}