Use type hints for resdata.well

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()

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)

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()