From 0abadf53faa5f59dcdb1f8adec41617109900e08 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BF=AB=E4=B9=90=E7=9A=84=E8=80=81=E9=BC=A0=E5=AE=9D?=
 =?UTF-8?q?=E5=AE=9D?= <42690037+LaoshuBaby@users.noreply.github.com>
Date: Thu, 7 Mar 2024 18:18:15 +0800
Subject: [PATCH] feat: try to find master_route colour for line
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

给地铁优化到这个程度已经很过分了！
---
 src/yuheng/plugin/viz_folium/__main__.py | 46 ++++++++++++++++++++++--
 1 file changed, 43 insertions(+), 3 deletions(-)

diff --git a/src/yuheng/plugin/viz_folium/__main__.py b/src/yuheng/plugin/viz_folium/__main__.py
index 7202f66..07fee75 100644
--- a/src/yuheng/plugin/viz_folium/__main__.py
+++ b/src/yuheng/plugin/viz_folium/__main__.py
@@ -1,7 +1,7 @@
 import os
 import sys
 import time
-from typing import List, Tuple, Union
+from typing import List, Tuple, Union, Optional
 
 import folium
 
@@ -205,8 +205,48 @@ def display(self, **kwargs) -> None:
                         work_burden_way_count += 1
                         time_this_way_start = time.time()
                         if kwargs.get("colour_original", False) == True:
-                            if obj.tags.get("colour", None) != None:
-                                colour = obj.tags.get("colour", None)
+
+                            def get_colour(
+                                target_way_id: Union[int, str],
+                                reference_carto: Carto,
+                            ) -> Optional[str]:
+                                if type(target_way_id) == type(str):
+                                    target_way_id = int(target_way_id)
+                                for (
+                                    way_id,
+                                    way_obj,
+                                ) in reference_carto.way_dict.items():
+                                    if (
+                                        way_obj.tags.get("colour", None)
+                                        != None
+                                    ):
+                                        return way_obj.tags["colour"]
+                                    else:
+                                        for (
+                                            relation_id,
+                                            relation_obj,
+                                        ) in (
+                                            reference_carto.relation_dict.items()
+                                        ):
+                                            for member in relation_obj.members:
+                                                if (
+                                                    member.type == "way"
+                                                    and member.ref
+                                                    == target_way_id
+                                                ):
+                                                    return (
+                                                        relation_obj.tags.get(
+                                                            "colour", None
+                                                        )
+                                                    )
+                                                else:
+                                                    # logger.debug(
+                                                    #     f"havent found any colour info for w{target_way_id}"
+                                                    # )
+                                                    return None
+
+                            if get_colour(id, element) != None:
+                                colour = get_colour(id, element)
                                 logger.info(
                                     f"Colour {colour} was detected in w{id}"
                                 )