mikebannis / rascontrol Goto Github PK

class River(object):
    def __init__(self, name, code, rc):
        self.name = name  # River name, string
        self.code = code  # River code, int - these start at 1, not 0
        self.rc = rc   # RasController object
        self.reaches = self._get_reaches()  # list of Reach objects
        self._update_reach_codes()

    # TODO -  the reach code should probably be pulled from the rascontller, although i+1 seems to work
    def _get_reaches(self):
        """
        Gets list of reaches for river represented by self
        :return: list of Reach objects
        """
        reaches = []
        reach_names = self.rc.geometry_getreaches(self.code)
        for i, name in enumerate(reach_names):
            new_reach = Reach(name, i+1, self)
            reaches.append(new_reach)
        return reaches

    def _update_reach_codes(self):
        """
        Sometimes the reach codes switch between the geometry and output files. This switches the codes after assigning
        each reach their respective nodes
        """
        
        if self.rc.output_getreaches(self.code)==self.rc.geometry_getreaches(self.code):
            pass
        else:
            new_reach_codes = {}
            new_reach_names = self.rc.output_getreaches(self.code)
            for code, reach in enumerate(new_reach_names):
                new_reach_codes[reach] = code+1
            
            for reach in self.reaches:
                updated_code = new_reach_codes[reach.name]
                reach.code = updated_code
        
    def __repr__(self):
        return 'River name = "'+self.name + '", River code = "' + str(self.code)+'"'

def geometry_getreaches(self, river_num):
        """
        Returns reach names in river numbered river_num using Geometry_GetReaches
        :param river_num: int
        :return: list of reach names
        """
        
        _, _, reaches = self.com_rc.Geometry_GetReaches(river_num, None, None)
        return reaches
    
def output_getreaches(self, river_num):
        """
        Returns reach names in river numbered river_num using Output_GetReaches
        :param river_num: int
        :return: list of reach names
        """
        
        _, _, reaches = self.com_rc.Output_GetReaches(river_num, None, None)
        return reaches