From f6430ae6c9664c1b8e25e95756ea9d74ae12e495 Mon Sep 17 00:00:00 2001 From: Eivind Jahren Date: Mon, 2 Mar 2026 13:34:19 +0100 Subject: [PATCH] Use type hints for resdata.well --- python/resdata/well/well_connection.py | 48 ++++++++----------- python/resdata/well/well_info.py | 42 +++++++---------- python/resdata/well/well_segment.py | 39 ++++++---------- python/resdata/well/well_state.py | 65 +++++++++++--------------- python/resdata/well/well_time_line.py | 16 ++----- 5 files changed, 82 insertions(+), 128 deletions(-) diff --git a/python/resdata/well/well_connection.py b/python/resdata/well/well_connection.py index 005f07ea5..d5b9c341d 100644 --- a/python/resdata/well/well_connection.py +++ b/python/resdata/well/well_connection.py @@ -42,54 +42,46 @@ class WellConnection(BaseCClass): def __init__(self): raise NotImplementedError("Class can not be instantiated directly") - def isOpen(self): - """@rtype: bool""" + def isOpen(self) -> bool: return self._is_open() - def ijk(self): - """@rtype: tuple of (int, int, int)""" + def ijk(self) -> tuple[int, int, int]: i = self._i() j = self._j() k = self._k() return i, j, k - def direction(self): - """@rtype: WellConnectionDirection""" + def direction(self) -> WellConnectionDirection: return self._get_dir() - def segmentId(self): - """@rtype: int""" + def segmentId(self) -> int: return self._segment_id() - def isFractureConnection(self): - """@rtype: bool""" + def isFractureConnection(self) -> bool: return self._fracture_connection() - def isMatrixConnection(self): - """@rtype: bool""" + def isMatrixConnection(self) -> bool: return self._matrix_connection() - def connectionFactor(self): - """@rtype: float""" + def connectionFactor(self) -> float: return self._connection_factor() - def __eq__(self, other): + def __eq__(self, other) -> bool: return self._equal(other) - def __hash__(self): + def __hash__(self) -> int: return id(self) - def __ne__(self, other): + def __ne__(self, other) -> bool: return not self == other def free(self): pass - def isMultiSegmentWell(self): - """@rtype: bool""" + def isMultiSegmentWell(self) -> bool: return self._is_msw() - def __repr__(self): + def __repr__(self) -> str: ijk = str(self.ijk()) frac = "fracture " if self.isFractureConnection() else "" open_ = "open " if self.isOpen() else "shut " @@ -111,26 +103,26 @@ def __repr__(self): ) ) - def gasRate(self): + def gasRate(self) -> float: return self._gas_rate() - def waterRate(self): + def waterRate(self) -> float: return self._water_rate() - def oilRate(self): + def oilRate(self) -> float: return self._oil_rate() - def volumeRate(self): + def volumeRate(self) -> float: return self._volume_rate() - def gasRateSI(self): + def gasRateSI(self) -> float: return self._gas_rate_si() - def waterRateSI(self): + def waterRateSI(self) -> float: return self._water_rate_si() - def oilRateSI(self): + def oilRateSI(self) -> float: return self._oil_rate_si() - def volumeRateSI(self): + def volumeRateSI(self) -> float: return self._volume_rate_si() diff --git a/python/resdata/well/well_info.py b/python/resdata/well/well_info.py index d86e9e4f3..51b9243da 100644 --- a/python/resdata/well/well_info.py +++ b/python/resdata/well/well_info.py @@ -1,4 +1,6 @@ from os.path import isfile +from collections.abc import Iterator + from cwrap import BaseCClass from resdata.grid import Grid from resdata.resfile.rd_file import ResdataFile @@ -26,11 +28,12 @@ class WellInfo(BaseCClass): "rd_well_time_line_ref well_info_get_ts(rd_well_info, char*)" ) - def __init__(self, grid, rst_file=None, load_segment_information=True): - """ - @type grid: Grid - @type rst_file: str or ResdataFile or list of str or list of ResdataFile - """ + def __init__( + self, + grid: Grid, + rst_file: None | list[str | ResdataFile] | str | ResdataFile = None, + load_segment_information: bool = True, + ): c_ptr = self._alloc(grid) super(WellInfo, self).__init__(c_ptr) if not c_ptr: @@ -43,19 +46,13 @@ def __init__(self, grid, rst_file=None, load_segment_information=True): else: self.addWellFile(rst_file, load_segment_information) - def __repr__(self): + def __repr__(self) -> str: return "WellInfo(well_count = %d) at 0x%x" % (len(self), self._address()) - def __len__(self): - """@rtype: int""" + def __len__(self) -> int: return self._get_well_count() - def __getitem__(self, item): - """ - @type item: int or str - @rtype: WellTimeLine - """ - + def __getitem__(self, item: int | str) -> WellTimeLine: if isinstance(item, str): if not item in self: raise KeyError("The well '%s' is not in this set." % item) @@ -70,31 +67,26 @@ def __getitem__(self, item): return self._get_ts(well_name).setParent(self) - def __iter__(self): - """@rtype: iterator of WellTimeLine""" + def __iter__(self) -> Iterator[WellTimeLine]: index = 0 while index < len(self): yield self[index] index += 1 - def allWellNames(self): - """@rtype: list of str""" + def allWellNames(self) -> list[str]: return [self._iget_well_name(index) for index in range(0, len(self))] - def __contains__(self, item): - """ - @type item: str - @rtype: bool - """ + def __contains__(self, item: str) -> bool: return self._has_well(item) def _assert_file_exists(self, rst_file): if not isfile(rst_file): raise IOError("No such file %s" % rst_file) - def addWellFile(self, rst_file, load_segment_information): - """@type rstfile: str or ResdataFile""" + def addWellFile( + self, rst_file: str | ResdataFile, load_segment_information: bool + ) -> None: if isinstance(rst_file, str): self._assert_file_exists(rst_file) self._load_rstfile(rst_file, load_segment_information) diff --git a/python/resdata/well/well_segment.py b/python/resdata/well/well_segment.py index 7c4b7e684..edf8e8691 100644 --- a/python/resdata/well/well_segment.py +++ b/python/resdata/well/well_segment.py @@ -24,62 +24,51 @@ class WellSegment(BaseCClass): def __init__(self): raise NotImplementedError("Class can not be instantiated directly") - def free(self): + def free(self) -> None: pass - def __repr__(self): + def __repr__(self) -> str: return "WellSegment(%s) at 0x%x" % (str(self), self._address()) - def __str__(self): + def __str__(self) -> str: return "{Segment ID:%d BranchID:%d Length:%g}" % ( self.id(), self.branchId(), self.length(), ) - def id(self): - """@rtype: int""" + def id(self) -> int: return self._id() - def linkCount(self): - """@rtype: int""" + def linkCount(self) -> int: return self._link_count() - def branchId(self): - """@rtype: int""" + def branchId(self) -> int: return self._branch_id() - def outletId(self): - """@rtype: int""" + def outletId(self) -> int: return self._outlet_id() - def isActive(self): - """@rtype: bool""" + def isActive(self) -> bool: return self._active() - def isMainStem(self): - """@rtype: bool""" + def isMainStem(self) -> bool: return self._main_stem() - def isNearestWellHead(self): - """@rtype: bool""" + def isNearestWellHead(self) -> bool: return self._nearest_wellhead() - def depth(self): - """@rtype: float""" + def depth(self) -> float: return self._depth() def __len__(self): return self.length() - def length(self): - """@rtype: float""" + def length(self) -> float: return self._length() - def totalLength(self): - """@rtype: float""" + def totalLength(self) -> float: return self._total_length() - def diameter(self): - """@rtype: float""" + def diameter(self) -> float: return self._diameter() diff --git a/python/resdata/well/well_state.py b/python/resdata/well/well_state.py index f0962bd3d..fc9788cfa 100644 --- a/python/resdata/well/well_state.py +++ b/python/resdata/well/well_state.py @@ -1,7 +1,7 @@ from cwrap import BaseCClass from resdata import ResdataPrototype -from resdata.well import WellType, WellConnection +from resdata.well import WellType, WellConnection, WellSegment from resdata.util.util import CTime @@ -62,44 +62,36 @@ class WellState(BaseCClass): def __init__(self): raise NotImplementedError("Class can not be instantiated directly") - def name(self): - """@rtype: str""" + def name(self) -> str: return self._get_name() - def isOpen(self): - """@rtype: bool""" + def isOpen(self) -> bool: return self._is_open() - def free(self): + def free(self) -> None: pass - def wellHead(self): + def wellHead(self) -> WellConnection: well_head = self._get_global_well_head() well_head.setParent(self) return well_head - def wellNumber(self): - """@rtype: int""" + def wellNumber(self) -> int: return self._well_number() - def reportNumber(self): - """@rtype: int""" + def reportNumber(self) -> int: return self._report_number() - def simulationTime(self): - """@rtype: CTime""" + def simulationTime(self) -> CTime: return self._sim_time() - def wellType(self): - """@rtype: WellType""" + def wellType(self) -> WellType: return self._well_type() - def hasGlobalConnections(self): - """@rtype: bool""" + def hasGlobalConnections(self) -> bool: return self._has_global_connections() - def globalConnections(self): - """@rtype: list of WellConnection""" + def globalConnections(self) -> list[WellConnection]: global_connections = self._get_global_connections() if global_connections is None: return [] @@ -113,20 +105,18 @@ def globalConnections(self): values.append(value) return values - def __len__(self): + def __len__(self) -> int: return self.numSegments() def __getitem__(self, idx): return self.igetSegment(idx) - def numSegments(self): - """@rtype: int""" + def numSegments(self) -> int: segment_collection = self._get_segment_collection() count = self._segment_collection_size(segment_collection) return count - def segments(self): - """@rtype: list of WellSegment""" + def segments(self) -> list[WellSegment]: segment_collection = self._get_segment_collection() values = [] @@ -138,8 +128,7 @@ def segments(self): return values - def igetSegment(self, seg_idx): - """@rtype: WellSegment""" + def igetSegment(self, seg_idx) -> WellSegment: if seg_idx < 0: seg_idx += len(self) @@ -153,15 +142,13 @@ def igetSegment(self, seg_idx): self ) - def isMultiSegmentWell(self): - """@rtype: bool""" + def isMultiSegmentWell(self) -> bool: return self._is_msw() - def hasSegmentData(self): - """@rtype: bool""" + def hasSegmentData(self) -> bool: return self._has_segment_data() - def __repr__(self): + def __repr__(self) -> str: name = self.name() if name: name = "%s" % name @@ -180,26 +167,26 @@ def __repr__(self): ) return self._create_repr(cnt) - def gasRate(self): + def gasRate(self) -> float: return self._gas_rate() - def waterRate(self): + def waterRate(self) -> float: return self._water_rate() - def oilRate(self): + def oilRate(self) -> float: return self._oil_rate() - def volumeRate(self): + def volumeRate(self) -> float: return self._volume_rate() - def gasRateSI(self): + def gasRateSI(self) -> float: return self._gas_rate_si() - def waterRateSI(self): + def waterRateSI(self) -> float: return self._water_rate_si() - def oilRateSI(self): + def oilRateSI(self) -> float: return self._oil_rate_si() - def volumeRateSI(self): + def volumeRateSI(self) -> float: return self._volume_rate_si() diff --git a/python/resdata/well/well_time_line.py b/python/resdata/well/well_time_line.py index f77f2b741..6db3be98d 100644 --- a/python/resdata/well/well_time_line.py +++ b/python/resdata/well/well_time_line.py @@ -14,19 +14,13 @@ class WellTimeLine(BaseCClass): def __init__(self): raise NotImplementedError("Class can not be instantiated directly") - def getName(self): + def getName(self) -> str: return self._name() - def __len__(self): - """@rtype: int""" + def __len__(self) -> int: return self._size() - def __getitem__(self, index): - """ - @type index: int - @rtype: WellState - """ - + def __getitem__(self, index: int) -> WellState: if index < 0: index += len(self) @@ -35,10 +29,10 @@ def __getitem__(self, index): return self._iget(index).setParent(self) - def free(self): + def free(self) -> None: pass - def __repr__(self): + def __repr__(self) -> str: n = self.getName() l = len(self) cnt = "name = %s, size = %d" % (n, l)