Provide an expand_grid function like in R
Change-Id: Ice4849b7ab35243066fb7d827f0ac35110ff211b
diff --git a/KorAPClient/__init__.py b/KorAPClient/__init__.py
index 77bb10c..14ca5a8 100644
--- a/KorAPClient/__init__.py
+++ b/KorAPClient/__init__.py
@@ -1,7 +1,9 @@
__pdoc__ = {'tests': False}
import warnings
+from itertools import product
+import pandas as pd
import rpy2.robjects as robjects
import rpy2.robjects.packages as packages
import rpy2.robjects.pandas2ri as pandas2ri
@@ -26,6 +28,27 @@
robjects.conversion.set_conversion(robjects.default_converter + pandas2ri.converter + korapclient_converter)
+def expand_grid(dictionary):
+ """Create a oandas DataFrame from all combinations of inputs
+
+ - **dictionary** - dict with variable names as keys and their values as vectors
+
+ Returns:
+ DataFrame with column names as specified by the dictionary key and all combinations of the specified values
+ in the rows.
+
+ Example:
+ ```
+ $ df = expand_grid({"Year": range(2010, 2019), "Country": ["DE", "CH"] })
+
+ $ df["vc"] = "textType=/Zeit.*/ & pubPlaceKey = " + df.Country + " & pubDate in " + list(map(str, df.Year))
+ ```
+ """
+
+ return pd.DataFrame([row for row in product(*dictionary.values())],
+ columns=dictionary.keys())
+
+
# noinspection PyPep8Naming
class KorAPConnection(RS4):
"""Connection to a KorAP server."""
diff --git a/examples/hello_world_interactive.py b/examples/hello_world_interactive.py
index f93e799..9aaa6a2 100755
--- a/examples/hello_world_interactive.py
+++ b/examples/hello_world_interactive.py
@@ -1,14 +1,10 @@
#!/usr/bin/env python3
import altair as alt
-import pandas as pd
-from KorAPClient import KorAPClient, KorAPConnection
+from KorAPClient import KorAPClient, KorAPConnection, expand_grid
QUERY = "Hello World"
-YEARS = range(2010, 2019)
-COUNTRIES = ["DE", "CH"]
-
-df = pd.DataFrame(YEARS, columns=["Year"], dtype=str).merge(pd.DataFrame(COUNTRIES, columns=["Country"]), how="cross")
-df["vc"] = "textType=/Zeit.*/ & pubPlaceKey = " + df.Country + " & pubDate in " + df.Year
+df = expand_grid({"Year": range(2010, 2019), "Country": ["DE", "CH"]})
+df["vc"] = "textType=/Zeit.*/ & pubPlaceKey = " + df.Country + " & pubDate in " + list(map(str, df.Year))
kcon = KorAPConnection(verbose=True)