speed up dry run

Author: gsheni

pymetrics/main.py

@@ -51,7 +51,9 @@ def collect_pypi_downloads(
     LOGGER.info(f'Collecting new downloads for projects={projects}')
 
     csv_path = get_path(output_folder, 'pypi.csv')
-    previous = get_previous_pypi_downloads(input_file=None, output_folder=output_folder)
+    previous = get_previous_pypi_downloads(
+        input_file=None, output_folder=output_folder, dry_run=dry_run
+    )
 
     pypi_downloads = get_pypi_downloads(
         projects=projects,

pymetrics/pypi.py

@@ -129,7 +129,6 @@ def get_pypi_downloads(
     if previous is not None:
         if isinstance(projects, str):
             projects = (projects,)
-
         previous_projects = previous[previous.project.isin(projects)]
         min_date = previous_projects.timestamp.min().date()
         max_date = previous_projects.timestamp.max().date()

pymetrics/summarize.py

@@ -2,6 +2,7 @@
 
 import logging
 import os
+import random
 
 import pandas as pd
 from packaging.version import Version, parse
@@ -110,7 +111,7 @@ def _sum_counts(base_count, dep_to_count, parent_to_count):
     return base_count + sum(parent_to_count.values()) + sum(dep_to_count.values())
 
 
-def get_previous_pypi_downloads(input_file, output_folder):
+def get_previous_pypi_downloads(input_file, output_folder, dry_run=False):
     """Read pypi.csv and return a DataFrame of the downloads.
 
     Args:
@@ -138,12 +139,12 @@ def get_previous_pypi_downloads(input_file, output_folder):
             'cpu': pd.CategoricalDtype(),
         },
     }
+    if dry_run:
+        read_csv_kwargs['nrows'] = 10_000
     data = load_csv(csv_path, read_csv_kwargs=read_csv_kwargs)
     LOGGER.info('Parsing version column to Version class objects')
-    if data and 'version' in data.columns:
+    if 'version' in data.columns:
         data['version'] = data['version'].apply(parse)
-    if not data:
-        data = pd.DataFrame()
     return data