added add_map_legend example

Author: mxaviersmp

pymove/utils/visual.py

@@ -40,6 +40,69 @@ def add_map_legend(m: Map, title: Text, items: List[Tuple]):
     ----------
     https://github.com/python-visualization/folium/issues/528#issuecomment-421445303
 
+    Examples
+    --------
+    >>> import folium
+    >>> from pymove.utils.visual import add_map_legend
+    >>> df
+              lat          lon              datetime  id
+    0   39.984094   116.319236   2008-10-23 05:53:05   1
+    1   39.984198   116.319322   2008-10-23 05:53:06   1
+    2   39.984224   116.319402   2008-10-23 05:53:11   1
+    3   39.984211   116.319389   2008-10-23 05:53:16   2
+    4   39.984217   116.319422   2008-10-23 05:53:21   2
+    >>> m = folium.Map(location=[df.lat.median(), df.lon.median()])
+    >>> folium.PolyLine(mdf[['lat', 'lon']], color='red').add_to(m)
+    >>> pm.visual.add_map_legend(m, 'Color by ID', [(1, 'red')])
+    >>> m.get_root().to_dict()
+    {
+        "name": "Figure",
+        "id": "1d32230cd6c54b19b35ceaa864e61168",
+        "children": {
+            "map_6f1abc8eacee41e8aa9d163e6bbb295f": {
+                "name": "Map",
+                "id": "6f1abc8eacee41e8aa9d163e6bbb295f",
+                "children": {
+                    "openstreetmap": {
+                        "name": "TileLayer",
+                        "id": "f58c3659fea348cb828775f223e1e6a4",
+                        "children": {}
+                    },
+                    "poly_line_75023fd7df01475ea5e5606ddd7f4dd2": {
+                        "name": "PolyLine",
+                        "id": "75023fd7df01475ea5e5606ddd7f4dd2",
+                        "children": {}
+                    }
+                }
+            },
+            "map_legend": {  # legend element
+                "name": "MacroElement",
+                "id": "72911b4418a94358ba8790aab93573d1",
+                "children": {}
+            }
+        },
+        "header": {
+            "name": "Element",
+            "id": "e46930fc4152431090b112424b5beb6a",
+            "children": {
+                "meta_http": {
+                    "name": "Element",
+                    "id": "868e20baf5744e82baf8f13a06849ecc",
+                    "children": {}
+                }
+            }
+        },
+        "html": {
+            "name": "Element",
+            "id": "9c4da9e0aac349f594e2d23298bac171",
+            "children": {}
+        },
+        "script": {
+            "name": "Element",
+            "id": "d092078607c04076bf58bd4593fa1684",
+            "children": {}
+        }
+    }
     """
     item = "<li><span style='background:%s;'></span>%s</li>"
     list_items = '\n'.join([item % (c, n) for (n, c) in items])
@@ -136,7 +199,7 @@ def add_map_legend(m: Map, title: Text, items: List[Tuple]):
     macro = MacroElement()
     macro._template = Template(template)
 
-    m.get_root().add_child(macro)
+    m.get_root().add_child(macro, name='map_legend')
 
 
 def generate_color() -> Text:
@@ -232,11 +295,11 @@ def cmap_hex_color(cmap: ListedColormap, i: int) -> Text:
     Examples
     --------
     >>> from pymove.utils.visual import  cmap_hex_color
-    >>> # import matplotlib.pyplot as plt
-    >>> # jet = plt.get_cmap('jet') // This comand generates a Linear Segmented Colormap
-    >>> print(cmap_hex_color(jet,0))
+    >>> import matplotlib.pyplot as plt
+    >>> jet = plt.get_cmap('jet')  # This comand generates a Linear Segmented Colormap
+    >>> print(cmap_hex_color(jet, 0))
     '#000080'
-    >>> print(cmap_hex_color(jet,1))
+    >>> print(cmap_hex_color(jet, 1))
     '#000084'
     """
     return rgb2hex(cmap(i))
@@ -259,9 +322,8 @@ def get_cmap(cmap: Text) -> Colormap:
     Examples
     --------
     >>> from pymove.utils.visual import  get_cmap
-    >>> print(get_cmap('Greys'), type(get_cmap('Greys')))
+    >>> print(get_cmap('Greys')
     <matplotlib.colors.LinearSegmentedColormap object at 0x7f743fc04bb0>
-     <class 'matplotlib.colors.LinearSegmentedColormap'>
     """
     return _get_cmap(cmap)
 
@@ -293,11 +355,16 @@ def save_wkt(
     0   39.984094   116.319236   2008-10-23 05:53:05   1
     1   39.984198   116.319322   2008-10-23 05:53:06   1
     2   39.984224   116.319402   2008-10-23 05:53:11   1
-    3   39.984211   116.319389   2008-10-23 05:53:16   1
-    4   39.984217   116.319422   2008-10-23 05:53:21   1
+    3   39.984211   116.319389   2008-10-23 05:53:16   2
+    4   39.984217   116.319422   2008-10-23 05:53:21   2
     >>> save_wkt(df, 'test.wkt', 'id')
+    >>> with open('test.wtk') as f:
+    >>>     print(f.read())
+    'id;linestring'
+    '1;LINESTRING(116.319236 39.984094,116.319322 39.984198,116.319402 39.984224)'
+    '2;LINESTRING(116.319389 39.984211,116.319422 39.984217)'
     """
-    str_ = '%s;linestring\n' % label_id
+    wtk = '%s;linestring\n' % label_id
     ids = move_data[label_id].unique()
     for id_ in ids:
         move_df = move_data[move_data[label_id] == id_]
@@ -307,6 +374,6 @@ def save_wkt(
             for x in move_df[[LONGITUDE, LATITUDE]].values
         )
         curr_str += ')\n'
-        str_ += curr_str
+        wtk += curr_str
     with open(filename, 'w') as f:
-        f.write(str_)
+        f.write(wtk)