0.4.0 add Nether rails + day mode

CHANGELOG.md

@@ -1,3 +1,8 @@
+# 0.4.0
+
++ Added nether lines
++ Added "day mode" colors
+
 # 0.3.0
 
 + Added colors

src/micorail/__init__.py

@@ -6,7 +6,7 @@ from .utils import HourMin
 from .rails import main_rails
 from .triwat import main_triwat
 
-__version__ = "0.3.0"
+__version__ = "0.4.0"
 
 def build_parser():
     parser = argparse.ArgumentParser()
@@ -18,8 +18,10 @@ def build_parser():
     parser_r.add_argument('start', help='station of start', default=None, nargs='?')
     parser_r.add_argument('end', help='station of end', default=None, nargs='?')
     parser_r.add_argument('-t', '--time', help='time of start', type=HourMin, default=HourMin(0))
+    parser_r.add_argument('-n', '--nether', help='use the Nether', action='store_true')
     parser_r.add_argument('--legacy', action='store_true', help="use legacy network")
     parser_r.add_argument('-q', '--search', help='search a station by its name')
+    parser_r.add_argument('--day', help='day mode', action='store_true')
 
     parser_t = parsers.add_parser('triwat', aliases=('t',))
     parser_t.add_argument('--multi', help="multiple posts", action='store_true')

src/micorail/rails.py

@@ -7,7 +7,7 @@ from math import ceil
 from heapq import heapify, heappush, heappop
 from suou import chalk
 
-from .utils import HourMin
+from .utils import HourMin, take_first
 
 ALL_DATA = json.load(open('data/network.json'))
 LEGACY_DATA = json.load(open('data/network.1.json'))
@@ -40,19 +40,18 @@ class RouteStep:
     line: str
 
 def find_station(code: str) -> Station | None:
-    st_name = ALL_DATA.get('stations', {}).get(code)
+    if not code:
+        return
+    if code.startswith('N-'):
+        st_name = take_first(ALL_DATA.get('stations', {}).get(code[2:], code[2:])) + ' Nether'
+    else:
+        st_name = take_first(ALL_DATA.get('stations', {}).get(code, code))
     if st_name:
         return Station(
             name = st_name,
             code = code
         )
 
-def take_first(s):
-    if isinstance(s, (str, bytes)):
-        return s
-    elif hasattr(s, '__iter__'):
-        return list(s)[0]
-    return s
 
 def build_route_list(line_stops: list, line_code):
     route_list = []
@@ -76,9 +75,9 @@ def build_route_list(line_stops: list, line_code):
         last_step = cur_step
     return route_list
 
-def build_all_lines():
+def build_all_lines(dim = 'overworld'):
     lines = {}
-    for line_data in ALL_DATA['lines']['overworld']:
+    for line_data in ALL_DATA['lines'][dim]:
         
         lines[line_data['code']] = Line(
             code = line_data['code'],
@@ -88,12 +87,13 @@ def build_all_lines():
     return lines
 
 ALL_LINES = build_all_lines()
+NETHER_LINES = build_all_lines('the_nether')
 
 ## TODO algorithms of research
-def find_route(start: str, stop: str):
+def find_route(start: str, stop: str, nether = False):
     steps_i = []
 
-    dist, prev = dijkstra(start)
+    dist, prev = dijkstra(start, nether=nether)
 
     cur = stop
     steps_i.append((find_station(stop), dist[stop]))
@@ -105,8 +105,10 @@ def find_route(start: str, stop: str):
     
     return steps_i
 
+NETHER_DELAY = 2
+
 @lru_cache()
-def find_neighbors(start):
+def find_neighbors(start: str, *, nether = False):
     neighs = []
 
     for line in ALL_LINES.values():
@@ -118,10 +120,28 @@ def find_neighbors(start):
                     neighs.append((rs.target.code, rs.time))
                 if rs.target.code == start:
                     neighs.append((rs.origin.code, rs.time))
+    if nether:
+        if start.startswith('N-'):
+            for line in NETHER_LINES.values():
+                line: Line
+                if start[2:] in line:
+                    for rs in line.route:
+                        rs: RouteStep
+                        if rs.origin.code == start[2:]:
+                            neighs.append(('N-' + rs.target.code, rs.time))
+                        if rs.target.code == start[2:]:
+                            neighs.append(('N-' + rs.origin.code, rs.time))
+            neighs.append((start[2:], NETHER_DELAY))
+        else:
+            for line in NETHER_LINES.values():
+                line: Line
+                if start in line:
+                    neighs.append((f'N-{start}', NETHER_DELAY))
+                    break
     
     return neighs
 
-def dijkstra(start: str):
+def dijkstra(start: str, *, nether: bool = False):
     dist = {node: INFINITY for node in ALL_DATA['stations']}
     dist[start] = 0
     prev = {node: None for node in ALL_DATA['stations']}
@@ -138,7 +158,7 @@ def dijkstra(start: str):
             continue
         visited.add(cur_node)
 
-        for neigh, time in find_neighbors(cur_node):
+        for neigh, time in find_neighbors(cur_node, nether = nether):
             tentative_dist = cur_dist + time
             if tentative_dist < dist.setdefault(neigh, INFINITY):
                 dist[neigh] = tentative_dist
@@ -147,6 +167,14 @@ def dijkstra(start: str):
 
     return dist, prev
 
+def hm_color(hm: HourMin):
+    if hm < 6 * 60 or hm >= 20 * 60:
+        return chalk.blue
+    if hm < 10 * 60:
+        return chalk.cyan
+    if hm < 16 * 60:
+        return chalk.yellow
+    return chalk.purple
 
 
 def main_rails(args):
@@ -170,7 +198,16 @@ def main_rails(args):
         print('**', f'{chalk.red(f'Missing stations')}')
         return
 
-    route = find_route(st_start.code, st_end.code)
+    try:
+        route = find_route(st_start.code, st_end.code, nether=args.nether)
+    except KeyError as k:
+        print('**', f'{chalk.red(f'Route {k} is impassable')}')
+        return
 
     for st_step, time in route:
-        print(chalk.cyan(f'{HourMin(st_time + time)}'), chalk.bold(st_step.code), st_step.name)
+        hm_time = HourMin(st_time + time)
+        print(
+            (hm_color(hm_time) if args.day else chalk.yellow)(f'{hm_time}'), 
+            f'{chalk.red('N-')}{chalk.bold(st_step.code[2:])}' if st_step.code.startswith('N-') else chalk.bold(st_step.code),
+            st_step.name
+        )

src/micorail/utils.py

@@ -8,4 +8,11 @@ class HourMin(int):
         if len(args) == 1 and isinstance(args[0], str) and ':' in args[0]:
             h, m = args[0].split(':')
             return super().__new__(cls, int(h) * 60 + int(m))
-        return super().__new__(cls, *args)
+        return super().__new__(cls, *args)
+
+def take_first(s):
+    if isinstance(s, (str, bytes)):
+        return s
+    elif hasattr(s, '__iter__'):
+        return next(iter(s))
+    return s