__geo_interface__ support

wntr/network/base.py

@@ -2,31 +2,32 @@
 The wntr.network.base module includes base classes for network elements and 
 the network model.
 """
+from __future__ import annotations
+
 import logging
-import six
-from six import string_types
 import types
+from dataclasses import dataclass
 from wntr.utils.ordered_set import OrderedSet
+from typing import Literal
 
 import enum
-import sys
-from collections.abc import MutableSequence, MutableMapping
+from abc import ABC
+from collections.abc import  MutableMapping
 from collections import OrderedDict
-from wntr.utils.ordered_set import OrderedSet
 
 import abc
 
 logger = logging.getLogger(__name__)
 
 
-class AbstractModel(object):
+class AbstractModel():
     """
     Base class for water network models.
     """
     pass
 
 
-class Subject(object):
+class Subject():
     """
     Base class for the subject in an observer design pattern.
     """
@@ -44,7 +45,7 @@ def notify(self):
             o.update(self)
 
 
-class Observer(six.with_metaclass(abc.ABCMeta, object)):
+class Observer(ABC):
     """
     Base class for the observer in an observer design pattern.
     """
@@ -53,7 +54,60 @@ def update(self, subject):
         pass
 
 
-class Node(six.with_metaclass(abc.ABCMeta, object)):
+@dataclass
+class Geometry():
+    geometry_type: Literal['Point','LineString','Polygon']
+    coordinates: tuple[int,int] | list[tuple[int,int]] | list[list[tuple[int,int]]]
+
+    @property
+    def __geo_interface__(self):
+        return {
+            "type": self.geometry_type,
+            "coordinates": self.coordinates,
+        }
+
+
+class PhysicalElement(ABC):
+    def __init__(self, wn, name):
+        self._options = wn._options
+        self._node_reg = wn._node_reg
+        self._link_reg = wn._link_reg
+        self._controls = wn._controls
+        self._pattern_reg = wn._pattern_reg
+        self._curve_reg = wn._curve_reg
+
+        self._name = name
+
+    def __str__(self):
+        return self._name
+
+    @property
+    def name(self):
+        """str: The name of the node (read only)"""
+        return self._name
+
+    @property
+    @abc.abstractmethod
+    def geometry(self) -> Geometry:
+        """Contains the geometry of the object"""
+        pass
+
+    @abc.abstractmethod
+    def to_dict(self) -> dict:
+        return {}
+
+    @property
+    def __geo_interface__(self) -> dict:
+        return {
+            "type": "Feature",
+            "properties": self.to_dict(),
+            "geometry": self.geometry.__geo_interface__,
+        }
+
+    def to_ref(self):
+        return self._name
+
+class Node(PhysicalElement):
     """Base class for nodes.
 
     For details about the different subclasses, see one of the following:
@@ -101,7 +155,7 @@ class Node(six.with_metaclass(abc.ABCMeta, object)):
 
     """
     def __init__(self, wn, name):
-        self._name = name
+        super().__init__(wn, name)
         self._head = None
         self._demand = None
         self._pressure = None
@@ -113,13 +167,7 @@ def __init__(self, wn, name):
         self._leak_status = False
         self._leak_area = 0.0
         self._leak_discharge_coeff = 0.0
-        self._options = wn._options
-        self._node_reg = wn._node_reg
-        self._link_reg = wn._link_reg
-        self._controls = wn._controls
-        self._pattern_reg = wn._pattern_reg
-        self._curve_reg = wn._curve_reg
-        self._coordinates = [0,0]
+        self._coordinates = (0, 0)
         self._source = None
         self._is_isolated = False
 
@@ -145,9 +193,6 @@ def _compare(self, other):
                return True
         return False
 
-    def __str__(self):
-        return self._name
-
     def __repr__(self):
         return "<Node '{}'>".format(self._name)
 
@@ -219,11 +264,6 @@ def node_type(self):
         """str: The node type (read only)"""
         return 'Node'
 
-    @property
-    def name(self):
-        """str: The name of the node (read only)"""
-        return self._name
-
     @property
     def tag(self):
         """str: A tag or label for the node"""
@@ -255,16 +295,20 @@ def coordinates(self, coordinates):
             self._coordinates = tuple(coordinates)
         else:
             raise ValueError('coordinates must be a 2-tuple or len-2 list')
+
+    @property
+    def geometry(self) -> Geometry:
+        return Geometry("Point", self.coordinates)
 
     def to_dict(self):
         """Dictionary representation of the node"""
         d = {}
         d['name'] = self.name
         d['node_type'] = self.node_type
         for k in dir(self):
-            if not k.startswith('_') and \
+            if not k[0]==('_') and \
               k not in ['demand', 'head', 'leak_demand', 'leak_status', 
-                        'level', 'pressure', 'quality', 'vol_curve', 'head_timeseries']:
+                        'level', 'pressure', 'quality', 'vol_curve', 'head_timeseries', 'geometry']:
                 try:
                     val = getattr(self, k)
                     if not isinstance(val, types.MethodType):
@@ -281,11 +325,11 @@ def to_dict(self):
                 except DeprecationWarning: pass
         return d
 
-    def to_ref(self):
-        return self._name
 
 
-class Link(six.with_metaclass(abc.ABCMeta, object)):
+
+
+class Link(PhysicalElement):
     """Base class for links.
 
     For details about the different subclasses, see one of the following:
@@ -340,15 +384,8 @@ class Link(six.with_metaclass(abc.ABCMeta, object)):
 
     """
     def __init__(self, wn, link_name, start_node_name, end_node_name):
-        # Set the registries
-        self._options = wn._options
-        self._node_reg = wn._node_reg
-        self._link_reg = wn._link_reg
-        self._controls = wn._controls
-        self._pattern_reg = wn._pattern_reg
-        self._curve_reg = wn._curve_reg
-        # Set the link name
-        self._link_name = link_name
+        super().__init__(wn, link_name)
+
         # Set and register the starting node
         self._start_node = self._node_reg[start_node_name]
         self._node_reg.add_usage(start_node_name, (link_name, self.link_type))
@@ -401,11 +438,8 @@ def _compare(self, other):
             return False
         return True
 
-    def __str__(self):
-        return self._link_name
-
     def __repr__(self):
-        return "<Link '{}'>".format(self._link_name)
+        return "<Link '{}'>".format(self._name)
 
     @property
     def link_type(self):
@@ -439,8 +473,8 @@ def start_node(self):
         return self._start_node
     @start_node.setter
     def start_node(self, node):
-        self._node_reg.remove_usage(self.start_node_name, (self._link_name, self.link_type))
-        self._node_reg.add_usage(node.name, (self._link_name, self.link_type))
+        self._node_reg.remove_usage(self.start_node_name, (self._name, self.link_type))
+        self._node_reg.add_usage(node.name, (self._name, self.link_type))
         self._start_node = self._node_reg[node.name]
 
     @property
@@ -449,8 +483,8 @@ def end_node(self):
         return self._end_node
     @end_node.setter
     def end_node(self, node):
-        self._node_reg.remove_usage(self.end_node_name, (self._link_name, self.link_type))
-        self._node_reg.add_usage(node.name, (self._link_name, self.link_type))
+        self._node_reg.remove_usage(self.end_node_name, (self._name, self.link_type))
+        self._node_reg.add_usage(node.name, (self._name, self.link_type))
         self._end_node = self._node_reg[node.name]
 
     @property
@@ -463,10 +497,6 @@ def end_node_name(self):
         """str: The name of the end node (read only)"""
         return self._end_node.name
 
-    @property
-    def name(self):
-        """str: The link name (read-only)"""
-        return self._link_name
 
     @property
     def flow(self):
@@ -548,6 +578,10 @@ def vertices(self, points):
                 raise ValueError('vertices must be a list of 2-tuples')
         self._vertices = points
 
+    @property
+    def geometry(self) -> Geometry:
+        return Geometry("LineString",[self.start_node.coordinates, *self.vertices, self.end_node.coordinates])
+
     def to_dict(self):
         """Dictionary representation of the link"""
         d = {}
@@ -560,9 +594,9 @@ def to_dict(self):
         if hasattr(self, 'valve_type'):
             d['valve_type'] = self.valve_type
         for k in dir(self):
-            if not k.startswith('_') and k not in [
+            if not k[0]==('_') and k not in [
                 'flow', 'cv', 'friction_factor', 'headloss',
-                'quality', 'reaction_rate', 'setting', 'status', 'velocity', 'speed_timeseries',
+                'quality', 'reaction_rate', 'setting', 'status', 'velocity', 'speed_timeseries', 'geometry'
             ]:
                 val = getattr(self, k)
                 if not isinstance(val, types.MethodType):
@@ -580,8 +614,6 @@ def to_dict(self):
                         d[k] = val
         return d
 
-    def to_ref(self):
-        return self._name
 
 
 class Registry(MutableMapping):
@@ -623,7 +655,7 @@ def __getitem__(self, key):
 
 
     def __setitem__(self, key, value):
-        if not isinstance(key, string_types):
+        if not isinstance(key, str):
             raise ValueError('Registry keys must be strings')
         self._data[key] = value
 
@@ -641,8 +673,9 @@ def __delitem__(self, key):
             # Do not raise an exception if there is no key of that name
             return
 
-    def __iter__(self):
+    def __iter__(self):    
         return self._data.__iter__()
+
 
     def __len__(self):
         return len(self._data)
@@ -758,6 +791,13 @@ def to_list(self):
             l.append(v.to_dict())
         return l
 
+    @property
+    def __geo_interface__(self):
+        return {
+            'type': 'FeatureCollection',
+            'features': [v.__geo_interface__ for v in self._data.values()]
+        }
+
 
 class NodeType(enum.IntEnum):
     """

wntr/network/elements.py

@@ -908,7 +908,7 @@ def __init__(self, name, start_node_name, end_node_name, wn):
         self._reaction_rate = None
 
     def __repr__(self):
-        return "<Pipe '{}' from '{}' to '{}', length={}, diameter={}, roughness={}, minor_loss={}, check_valve={}, status={}>".format(self._link_name,
+        return "<Pipe '{}' from '{}' to '{}', length={}, diameter={}, roughness={}, minor_loss={}, check_valve={}, status={}>".format(self._name,
                        self.start_node, self.end_node, self.length, self.diameter, 
                        self.roughness, self.minor_loss, self.check_valve, str(self.status))
 
@@ -1302,7 +1302,7 @@ class HeadPump(Pump):
 #                                         # the _curve_coeffs were calculated
 
     def __repr__(self):
-        return "<Pump '{}' from '{}' to '{}', pump_type='{}', pump_curve={}, speed={}, status={}>".format(self._link_name,
+        return "<Pump '{}' from '{}' to '{}', pump_type='{}', pump_curve={}, speed={}, status={}>".format(self._name,
                    self.start_node, self.end_node, 'HEAD', self.pump_curve_name, 
                    self.speed_timeseries, str(self.status))
 
@@ -1325,8 +1325,8 @@ def pump_curve_name(self):
         return self._pump_curve_name
     @pump_curve_name.setter
     def pump_curve_name(self, name):
-        self._curve_reg.remove_usage(self._pump_curve_name, (self._link_name, 'Pump'))
-        self._curve_reg.add_usage(name, (self._link_name, 'Pump'))
+        self._curve_reg.remove_usage(self._pump_curve_name, (self._name, 'Pump'))
+        self._curve_reg.add_usage(name, (self._name, 'Pump'))
         self._curve_reg.set_curve_type(name, 'HEAD')
         self._pump_curve_name = name
         # delete the pump curve coefficients because they have to be recaulcated 
@@ -1515,7 +1515,7 @@ class PowerPump(Pump):
     """
 
     def __repr__(self):
-        return "<Pump '{}' from '{}' to '{}', pump_type='{}', power={}, speed={}, status={}>".format(self._link_name,
+        return "<Pump '{}' from '{}' to '{}', pump_type='{}', power={}, speed={}, status={}>".format(self._name,
                    self.start_node, self.end_node, 'POWER', self._base_power, 
                    self.speed_timeseries, str(self.status))
 
@@ -1538,7 +1538,7 @@ def power(self):
         return self._base_power
     @power.setter
     def power(self, value):
-        self._curve_reg.remove_usage(self._pump_curve_name, (self._link_name, 'Pump'))
+        self._curve_reg.remove_usage(self._pump_curve_name, (self._name, 'Pump'))
         self._base_power = value
 
 
@@ -1625,7 +1625,7 @@ def __init__(self, name, start_node_name, end_node_name, wn):
 
     def __repr__(self):
         fmt = "<Valve '{}' from '{}' to '{}', valve_type='{}', diameter={}, minor_loss={}, setting={}, status={}>"
-        return fmt.format(self._link_name,
+        return fmt.format(self._name,
                           self.start_node, self.end_node, self.__class__.__name__,
                           self.diameter, 
                           self.minor_loss, self.setting, str(self.status))
@@ -2037,8 +2037,8 @@ def headloss_curve_name(self):
         return self._headloss_curve_name
     @headloss_curve_name.setter
     def headloss_curve_name(self, name):
-        self._curve_reg.remove_usage(self._headloss_curve_name, (self._link_name, 'Valve'))
-        self._curve_reg.add_usage(name, (self._link_name, 'Valve'))
+        self._curve_reg.remove_usage(self._headloss_curve_name, (self._name, 'Valve'))
+        self._curve_reg.add_usage(name, (self._name, 'Valve'))
         self._curve_reg.set_curve_type(name, 'HEADLOSS')
         self._headloss_curve_name = name