Refactor Statistics class in statistics.py to utilize internal API for data fetching, streamline chart creation, and enhance heatmap generation methods. Remove unused imports and redundant code for improved clarity and maintainability.

Zingzy · Zingzy · commit 2df5ba588713 · 2025-05-23T04:31:00.000+05:30
diff --git a/py_spoo_url/_internal/__init__.py b/py_spoo_url/_internal/__init__.py
@@ -0,0 +1,9 @@
+"""
+Internal modules for py_spoo_url
+"""
+
+from .api import fetch_statistics
+from .plotting import make_chart, make_countries_heatmap, make_unique_countries_heatmap
+from .exporters import export_data
+
+__all__ = ["fetch_statistics", "make_chart", "make_countries_heatmap", "make_unique_countries_heatmap", "export_data"]
diff --git a/py_spoo_url/_internal/api.py b/py_spoo_url/_internal/api.py
@@ -0,0 +1,13 @@
+import requests
+import json
+from typing import Optional, Any
+
+
+def fetch_statistics(short_code: str, password: Optional[str] = None) -> Any:
+    url = f"https://spoo.me/stats/{short_code}"
+    params = {"password": password} if password else None
+    r = requests.post(url, data=params)
+    if r.status_code == 200:
+        return json.loads(r.text)
+    else:
+        raise Exception(f"Error {r.status_code}: {r.text}")
diff --git a/py_spoo_url/_internal/exporters.py b/py_spoo_url/_internal/exporters.py
@@ -0,0 +1,106 @@
+import pandas as pd
+import json
+import os
+import shutil
+import zipfile
+from typing import Literal
+from .utils import create_dataframes_from_data, create_general_info_dataframe, STANDARD_DATAFRAME_CONFIGS
+
+
+def export_to_excel(data, filename: str = "export.xlsx") -> None:
+    # Create all standard DataFrames using utility function
+    dataframes = create_dataframes_from_data(data, STANDARD_DATAFRAME_CONFIGS)
+    
+    # Create general info DataFrame
+    df_general_info = create_general_info_dataframe(data, "URL")
+    
+    # Excel sheet name mapping
+    sheet_mapping = {
+        "df_browser": "Browser",
+        "df_counter": "Counter", 
+        "df_country": "Country",
+        "df_os_name": "OS_Name",
+        "df_referrer": "Referrer",
+        "df_unique_browser": "Unique_Browser",
+        "df_unique_counter": "Unique_Counter",
+        "df_unique_country": "Unique_Country",
+        "df_unique_os_name": "Unique_OS_Name",
+        "df_unique_referrer": "Unique_Referrer",
+    }
+    
+    with pd.ExcelWriter(filename, engine="openpyxl") as writer:
+        # Write all standard DataFrames
+        for df_name, sheet_name in sheet_mapping.items():
+            if df_name in dataframes:
+                dataframes[df_name].to_excel(writer, sheet_name=sheet_name, index=False)
+        
+        # Write general info
+        df_general_info.to_excel(writer, sheet_name="General_Info", index=False)
+    
+    print(f"Data successfully written to {filename}")
+
+
+def export_to_csv(data, filename: str = "export.csv") -> None:
+    csv_directory = "csv_files"
+    os.makedirs(csv_directory, exist_ok=True)
+    
+    # Create all standard DataFrames using utility function
+    dataframes = create_dataframes_from_data(data, STANDARD_DATAFRAME_CONFIGS)
+    
+    # Create general info DataFrame (with empty string for URL column in CSV)
+    df_general_info = create_general_info_dataframe(data, "")
+    
+    # CSV file name mapping
+    csv_mapping = {
+        "df_browser": "browser.csv",
+        "df_counter": "counter.csv",
+        "df_country": "country.csv",
+        "df_os_name": "os_name.csv",
+        "df_referrer": "referrer.csv",
+        "df_unique_browser": "unique_browser.csv",
+        "df_unique_counter": "unique_counter.csv",
+        "df_unique_country": "unique_country.csv",
+        "df_unique_os_name": "unique_os_name.csv",
+        "df_unique_referrer": "unique_referrer.csv",
+    }
+    
+    # Save all DataFrames to CSV files
+    for df_name, csv_filename in csv_mapping.items():
+        if df_name in dataframes:
+            dataframes[df_name].to_csv(os.path.join(csv_directory, csv_filename), index=False)
+    
+    # Save general info
+    df_general_info.to_csv(os.path.join(csv_directory, "general_info.csv"), index=False)
+    
+    # Create zip file
+    with zipfile.ZipFile(f"{filename}.zip", "w") as zipf:
+        for root, dirs, files in os.walk(csv_directory):
+            for file in files:
+                file_path = os.path.join(root, file)
+                arcname = os.path.relpath(file_path, csv_directory)
+                zipf.write(file_path, arcname=arcname)
+    
+    # Clean up temporary directory
+    shutil.rmtree(csv_directory)
+    print(f"Data successfully written to {filename}.zip")
+
+
+def export_to_json(data, filename: str = "export.json") -> None:
+    with open(filename, "w") as w:
+        w.write(json.dumps(data, indent=4))
+    print(f"Data successfully written to {filename}")
+
+
+def export_data(
+    data,
+    filename: str = "export.xlsx",
+    filetype: Literal["csv", "xlsx", "json"] = "xlsx",
+) -> None:
+    if filetype == "xlsx":
+        export_to_excel(data, filename)
+    elif filetype == "json":
+        export_to_json(data, filename)
+    elif filetype == "csv":
+        export_to_csv(data, filename)
+    else:
+        raise ValueError("Invalid file type. Choose either 'csv', 'json' or 'xlsx'.")
diff --git a/py_spoo_url/_internal/plotting.py b/py_spoo_url/_internal/plotting.py
@@ -0,0 +1,108 @@
+import matplotlib.pyplot as plt  # type: ignore
+import matplotlib  # type: ignore
+from mpl_toolkits.axes_grid1 import make_axes_locatable  # type: ignore
+import geopandas as gpd  # type: ignore
+from typing import Literal, Dict
+
+
+def make_chart(
+    chart_data: Dict,
+    chart_type: Literal["bar", "pie", "line", "scatter", "hist", "box", "area"] = "bar",
+    data_label: str = None,
+    **kwargs,
+) -> plt.Figure:
+    matplotlib.rcParams["font.size"] = 15
+    matplotlib.rcParams["axes.labelcolor"] = "Black"
+
+    if chart_type == "bar":
+        plt.bar(chart_data.keys(), chart_data.values(), **kwargs)
+        if data_label in [
+            "last_n_days_analysis",
+            "clicks_analysis",
+            "unique_clicks_analysis",
+        ]:
+            plt.xticks(rotation=90)
+    elif chart_type == "pie":
+        plt.pie(chart_data.values(), labels=chart_data.keys(), **kwargs)
+    elif chart_type == "line":
+        plt.plot(chart_data.keys(), chart_data.values(), **kwargs)
+    elif chart_type == "scatter":
+        plt.scatter(chart_data.keys(), chart_data.values(), **kwargs)
+    elif chart_type == "hist":
+        plt.hist(list(chart_data.values()), **kwargs)
+    elif chart_type == "box":
+        plt.boxplot(list(chart_data.values()), **kwargs)
+    elif chart_type == "area":
+        plt.stackplot(chart_data.keys(), chart_data.values(), **kwargs)
+    else:
+        raise Exception(
+            "Invalid chart type. Valid chart types are: bar, pie, line, scatter, hist, box, area"
+        )
+    return plt
+
+
+def _create_heatmap(
+    data_analysis: Dict[str, int],
+    title: str,
+    merge_column: str = "NAME",
+    cmap: Literal["YlOrRd", "viridis", "plasma", "inferno", "RdPu_r"] = "YlOrRd",
+) -> plt.Figure:
+    matplotlib.rcParams["font.size"] = 15
+    matplotlib.rcParams["axes.labelcolor"] = "White"
+    world = gpd.read_file("py_spoo_url/data/ne_110m_admin_0_countries.zip")
+    world = world.merge(
+        gpd.GeoDataFrame(data_analysis.items(), columns=["Country", "Value"]),
+        how="left",
+        left_on=merge_column,
+        right_on="Country",
+    )
+    fig, ax = plt.subplots(
+        1, 1, figsize=(15, 10), facecolor=(32 / 255, 34 / 255, 37 / 255, 0.5)
+    )
+    plt.subplots_adjust(left=0.05, right=0.90, bottom=0.05, top=0.95)
+    for spine in ax.spines.values():
+        spine.set_color((46 / 255, 48 / 255, 53 / 255))
+        spine.set_linewidth(2)
+    ax.tick_params(labelcolor="white")
+    world.boundary.plot(ax=ax, linewidth=1)
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="5%", pad=0.1)
+    p = world.plot(
+        column="Value",
+        ax=ax,
+        legend=True,
+        cax=cax,
+        cmap=cmap,
+        edgecolor=None,
+        legend_kwds={"label": "Clicks"},
+        alpha=0.9,
+    )
+    p.set_facecolor((32 / 255, 34 / 255, 37 / 255, 0.5))
+    cbax = cax
+    cbax.tick_params(labelcolor="white")
+    plt.suptitle(title, x=0.5, y=0.95, fontsize=20, fontweight=3, color="white")
+    return plt
+
+
+def make_countries_heatmap(
+    country_analysis: Dict[str, int],
+    cmap: Literal["YlOrRd", "viridis", "plasma", "inferno", "RdPu_r"] = "YlOrRd",
+) -> plt.Figure:
+    return _create_heatmap(
+        data_analysis=country_analysis,
+        title="Countries Heatmap",
+        merge_column="NAME",
+        cmap=cmap,
+    )
+
+
+def make_unique_countries_heatmap(
+    unique_country_analysis: Dict[str, int],
+    cmap: Literal["YlOrRd", "viridis", "plasma", "inferno", "RdPu_r"] = "YlOrRd",
+) -> plt.Figure:
+    return _create_heatmap(
+        data_analysis=unique_country_analysis,
+        title="Unique Countries Heatmap",
+        merge_column="NAME",
+        cmap=cmap,
+    )
diff --git a/py_spoo_url/_internal/utils.py b/py_spoo_url/_internal/utils.py
@@ -0,0 +1,64 @@
+import pandas as pd
+from typing import Dict, List, Tuple
+
+
+def create_dataframes_from_data(data: Dict, dataframe_configs: List[Tuple[str, str, List[str]]]) -> Dict[str, pd.DataFrame]:
+    """
+    Create multiple DataFrames from data dictionary based on configuration.
+    
+    Args:
+        data: Raw data dictionary
+        dataframe_configs: List of tuples (data_key, df_name, column_names)
+    
+    Returns:
+        Dictionary mapping DataFrame names to DataFrames
+    """
+    dataframes = {}
+    for data_key, df_name, columns in dataframe_configs:
+        if data_key in data:
+            dataframes[df_name] = pd.DataFrame(data[data_key].items(), columns=columns)
+    return dataframes
+
+
+def create_general_info_dataframe(data: Dict, url_column_name: str = "URL") -> pd.DataFrame:
+    """
+    Create the general info DataFrame with standardized structure.
+    
+    Args:
+        data: Raw data dictionary
+        url_column_name: Name for the URL column (use "" for CSV export)
+    
+    Returns:
+        DataFrame with general information
+    """
+    return pd.DataFrame({
+        "TOTAL CLICKS": [data["total-clicks"]],
+        "TOTAL UNIQUE CLICKS": [data["total_unique_clicks"]],
+        url_column_name: [data["url"]],
+        "SHORT CODE": [data["_id"]],
+        "MAX CLICKS": [data["max-clicks"]],
+        "PASSWORD": [data["password"]],
+        "CREATION DATE": [data["creation-date"]],
+        "EXPIRED": [data["expired"]],
+        "AVERAGE DAILY CLICKS": [data["average_daily_clicks"]],
+        "AVERAGE MONTHLY CLICKS": [data["average_monthly_clicks"]],
+        "AVERAGE WEEKLY CLICKS": [data["average_weekly_clicks"]],
+        "LAST CLICK": [data["last-click"]],
+        "LAST CLICK BROSWER": [data["last-click-browser"]],
+        "LAST CLICK OS": [data["last-click-os"]],
+    })
+
+
+# Configuration for standard DataFrames
+STANDARD_DATAFRAME_CONFIGS = [
+    ("browser", "df_browser", ["Browser", "Count"]),
+    ("counter", "df_counter", ["Date", "Count"]),
+    ("country", "df_country", ["Country", "Count"]),
+    ("os_name", "df_os_name", ["OS_Name", "Count"]),
+    ("referrer", "df_referrer", ["Referrer", "Count"]),
+    ("unique_browser", "df_unique_browser", ["Browser", "Count"]),
+    ("unique_counter", "df_unique_counter", ["Date", "Count"]),
+    ("unique_country", "df_unique_country", ["Country", "Count"]),
+    ("unique_os_name", "df_unique_os_name", ["OS_Name", "Count"]),
+    ("unique_referrer", "df_unique_referrer", ["Referrer", "Count"]),
+] 
diff --git a/py_spoo_url/statistics.py b/py_spoo_url/statistics.py
