feat: class to manage base bcr data and some tests

Author: arcstur

src/graphs/bar_chart_race.py

@@ -1,6 +1,43 @@
 import pandas as pd
+from pandas.errors import ParserError
 import re
 
+class BaseDf:
+    def __init__(self, df):
+        self.df = df
+
+    def prepare(self):
+        self.verify_column_count()
+        self.prepare_date_column()
+        self.prepare_value_column()
+
+    def verify_column_count(self):
+        if not (3 <= self.df.shape[1] <= 6):
+            raise BaseDfException("number of columns must be between 3 and 6")
+        
+    def prepare_date_column(self):
+        original_name = self.df.columns[-1]
+        try:
+            self.df[original_name] = pd.to_datetime(self.df[original_name], format="ISO8601")
+        except (ValueError, ParserError):
+            raise BaseDfException("last column must be a date column")
+        self.df.rename(columns={original_name: "date"})
+
+    def prepare_value_column(self):
+        original_name = self.df.columns[-2]
+        try:
+            self.df[original_name] = self.df[original_name].astype("float")
+        except ValueError:
+            raise BaseDfException("second to last column must be a quantity column")
+        self.df.rename(columns={original_name: "value"})
+
+
+class BaseDfException(Exception):
+    def __init__(self, message):
+        self.message = message
+        return super().__init__(message)
+
+
 def process_bar_chart_race(df):
     """
     Process data for bar chart race visualization.
@@ -9,19 +46,11 @@ def process_bar_chart_race(df):
     :return: Processed data suitable for bar chart race or error message
     """
 
-    # RULES
-
-    # 1. Check if the number of columns is between 3 and 6
-    if not (3 <= df.shape[1] <= 6):
-        return {"failed": "data failed bar chart race rule - number of columns must be between 3 and 6"}
-
-    # 2. Check if the last column is a date column
-    if not df.iloc[:, -1].apply(is_datetime_string).any():
-        return {"failed": "data failed bar chart race rule - last column must be a date column"}
-
-    # 3. Check if the second to last column is a number-string column
-    if not df.iloc[:, -2].apply(lambda x: isinstance(x, str) and x.replace('.', '', 1).isdigit()).any():
-        return {"failed": "data failed bar chart race rule - second to last column must be a quantity column"}
+    bdf = BaseDf(df)
+    try:
+        bdf.prepare()
+    except BaseDfException as e:
+        return {"failed": e.message}
 
     # 4. Identify columns with word identifiers
     identifier_columns = identify_word_identifier_columns(df.iloc[:, :-2])

src/graphs/tests.py

@@ -1,7 +1,9 @@
 import pandas as pd
 from django.test import TestCase
+from django.utils.timezone import now
 
 from graphs.table import process_table
+from graphs.bar_chart_race import process_bar_chart_race
 
 
 class TestHelper:
@@ -39,3 +41,18 @@ def test_basic_table(self):
                 },
             ],
         )
+
+    def test_bcr_base_errors(self):
+        df = TestHelper.mock_df_table()
+        msg = process_bar_chart_race(df)["failed"]
+        self.assertEqual(msg, "number of columns must be between 3 and 6")
+        df["new column"] = "abcdef"
+        msg = process_bar_chart_race(df)["failed"]
+        self.assertEqual(msg, "last column must be a date column")
+        df["new column"] = now()
+        msg = process_bar_chart_race(df)["failed"]
+        self.assertEqual(msg, "second to last column must be a quantity column")
+        df["new column"] = 2
+        df["last column"] = now()
+        res = process_bar_chart_race(df)
+        self.assertNotIn("failed", res)

src/query/sparql.py

@@ -10,28 +10,32 @@ def df_from_query(sparql_string):
     :param sparql_string: SPARQL query string
     :return: DataFrame containing the results
     """
-    url = "https://query.wikidata.org/sparql"
-    params = {
-        "query": sparql_string,
-        "format": "json"
-    }
-    headers = {'User-agent': 'Wiki-Infographics 1.0'}
-    response = requests.get(url, headers=headers, params=params)
+    response = get_response(sparql_string)
 
     if response.status_code != 200:
         error_response = extract_error_message(response.text)
         return {"error": error_response}
 
     data = response.json()
-    results = data['results']['bindings']
-    variables = data['head']['vars']
+    results = data["results"]["bindings"]
+    variables = data["head"]["vars"]
 
-    # Convert results to DataFrame
-    df = pd.DataFrame([{var: binding.get(var, {}).get('value', None) for var in variables} for binding in results])
+    df = pd.DataFrame(
+        [
+            {var: binding.get(var, {}).get("value", None) for var in variables}
+            for binding in results
+        ]
+    )
 
     return df
 
 
+def get_response(sparql_string):
+    url = "https://query.wikidata.org/sparql"
+    params = {"query": sparql_string, "format": "json"}
+    headers = {"User-agent": "Wiki-Infographics 1.0"}
+    return requests.get(url, headers=headers, params=params)
+
 
 def extract_error_message(error_str):
     """
@@ -43,10 +47,10 @@ def extract_error_message(error_str):
     Returns:
         str: Formatted error message in the form of "MalformedQueryException: [error details]."
     """
-    
+
     pattern = r"MalformedQueryException: [^.]*\."
     match = re.search(pattern, error_str)
-    
+
     if match:
         return match.group(0)
     else: