From 36f2ff63884a23bef5a2001a9253dfc39cc601f0 Mon Sep 17 00:00:00 2001
From: Frank Anema <33519926+Conengmo@users.noreply.github.com>
Date: Wed, 28 Feb 2024 16:20:48 +0100
Subject: [PATCH] realtime plugin: allow str for JsCode arg (#1862)

* realtime plugin: allow str for JsCode arg

* JsCode: don't inherit from str

* ruff

* add tests for JsCode

* run black
---
 folium/plugins/realtime.py | 17 +++++++++--------
 folium/utilities.py        |  7 +++++--
 tests/test_utilities.py    | 14 ++++++++++++++
 3 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/folium/plugins/realtime.py b/folium/plugins/realtime.py
index 429a40955..fe56de7c6 100644
--- a/folium/plugins/realtime.py
+++ b/folium/plugins/realtime.py
@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Union
 
 from branca.element import MacroElement
 from jinja2 import Template
@@ -27,11 +27,11 @@ class Realtime(JSCSSMixin, MacroElement):
         on the map and stopped when layer is removed from the map
     interval: int, default 60000
         Automatic update interval, in milliseconds
-    get_feature_id: JsCode, optional
+    get_feature_id: str or JsCode, optional
         A JS function with a geojson `feature` as parameter
         default returns `feature.properties.id`
         Function to get an identifier to uniquely identify a feature over time
-    update_feature: JsCode, optional
+    update_feature: str or JsCode, optional
         A JS function with a geojson `feature` as parameter
         Used to update an existing feature's layer;
         by default, points (markers) are updated, other layers are discarded
@@ -44,7 +44,8 @@ class Realtime(JSCSSMixin, MacroElement):
 
 
     Other keyword arguments are passed to the GeoJson layer, so you can pass
-    `style`, `point_to_layer` and/or `on_each_feature`.
+    `style`, `point_to_layer` and/or `on_each_feature`. Make sure to wrap
+    Javascript functions in the JsCode class.
 
     Examples
     --------
@@ -95,8 +96,8 @@ def __init__(
         source: Union[str, dict, JsCode],
         start: bool = True,
         interval: int = 60000,
-        get_feature_id: Optional[JsCode] = None,
-        update_feature: Optional[JsCode] = None,
+        get_feature_id: Union[JsCode, str, None] = None,
+        update_feature: Union[JsCode, str, None] = None,
         remove_missing: bool = False,
         **kwargs
     ):
@@ -107,9 +108,9 @@ def __init__(
         kwargs["start"] = start
         kwargs["interval"] = interval
         if get_feature_id is not None:
-            kwargs["get_feature_id"] = get_feature_id
+            kwargs["get_feature_id"] = JsCode(get_feature_id)
         if update_feature is not None:
-            kwargs["update_feature"] = update_feature
+            kwargs["update_feature"] = JsCode(update_feature)
         kwargs["remove_missing"] = remove_missing
 
         # extract JsCode objects
diff --git a/folium/utilities.py b/folium/utilities.py
index 7e9416a40..98c757599 100644
--- a/folium/utilities.py
+++ b/folium/utilities.py
@@ -423,5 +423,8 @@ def get_and_assert_figure_root(obj: Element) -> Figure:
 class JsCode:
     """Wrapper around Javascript code."""
 
-    def __init__(self, js_code: str):
-        self.js_code = js_code
+    def __init__(self, js_code: Union[str, "JsCode"]):
+        if isinstance(js_code, JsCode):
+            self.js_code: str = js_code.js_code
+        else:
+            self.js_code = js_code
diff --git a/tests/test_utilities.py b/tests/test_utilities.py
index 54fc2837a..7d5c44765 100644
--- a/tests/test_utilities.py
+++ b/tests/test_utilities.py
@@ -4,6 +4,7 @@
 
 from folium import FeatureGroup, Map, Marker, Popup
 from folium.utilities import (
+    JsCode,
     _is_url,
     camelize,
     deep_copy,
@@ -216,3 +217,16 @@ def test_escape_double_quotes(text, result):
 )
 def test_javascript_identifier_path_to_array_notation(text, result):
     assert javascript_identifier_path_to_array_notation(text) == result
+
+
+def test_js_code_init_str():
+    js_code = JsCode("hi")
+    assert isinstance(js_code, JsCode)
+    assert isinstance(js_code.js_code, str)
+
+
+def test_js_code_init_js_code():
+    js_code = JsCode("hi")
+    js_code_2 = JsCode(js_code)
+    assert isinstance(js_code_2, JsCode)
+    assert isinstance(js_code_2.js_code, str)