Allow yyyy-mm for -sd & -ed, add --month shortcut (ofek#71)

Allow yyyy-mm for -sd & -ed, add --month shortcut

README.rst

@@ -307,6 +307,39 @@ Downloads between two YYYY-MM-DD dates
     | -------------- |
     |      9,572,911 |
 
+Downloads between two YYYY-MM dates
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+- A yyyy-mm ``--start-date`` defaults to the first day of the month
+- A yyyy-mm ``--end-date`` defaults to the last day of the month
+
+.. code-block:: console
+
+    $ pypinfo --start-date 2018-04 --end-date 2018-04 setuptools
+    Served from cache: True
+    Data processed: 0.00 B
+    Data billed: 0.00 B
+    Estimated cost: $0.00
+
+    | download_count |
+    | -------------- |
+    |      9,572,911 |
+
+Downloads for a single YYYY-MM month
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. code-block:: console
+
+    $ pypinfo --month 2018-04 setuptools
+    Served from cache: True
+    Data processed: 0.00 B
+    Data billed: 0.00 B
+    Estimated cost: $0.00
+
+    | download_count |
+    | -------------- |
+    |      9,572,911 |
+
 Percentage of Python 3 downloads of the top 100 projects in the past year
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
@@ -346,11 +379,22 @@ Changelog
 
 Important changes are emphasized.
 
+Unreleased
+^^^^^^^^^^
+
+- Allow yyyy-mm[-dd] ``--start-date`` and ``--end-date``:
+
+  - A yyyy-mm ``--start-date`` defaults to the first day of the month
+  - A yyyy-mm ``--end-date`` defaults to the last day of the month
+
+- Add ``--month`` as a shortcut to ``--start-date`` and ``--end-date``
+  for a single yyyy-mm month
+
 15.0.0
 ^^^^^^
 
 - Allow yyyy-mm-dd dates
-- Add --all option, default to only showing downloads via pip
+- Add ``--all`` option, default to only showing downloads via pip
 - Add download total row
 
 14.0.0

pypinfo/cli.py

@@ -10,6 +10,7 @@
     create_client,
     create_config,
     format_json,
+    month_ends,
     parse_query_result,
     tabulate,
 )
@@ -77,8 +78,9 @@
 @click.option('--timeout', '-t', type=int, default=120000, help='Milliseconds. Default: 120000 (2 minutes)')
 @click.option('--limit', '-l', help='Maximum number of query results. Default: 10')
 @click.option('--days', '-d', help='Number of days in the past to include. Default: 30')
-@click.option('--start-date', '-sd', help='Must be negative or YYYY-MM-DD. Default: -31')
-@click.option('--end-date', '-ed', help='Must be negative or YYYY-MM-DD. Default: -1')
+@click.option('--start-date', '-sd', help='Must be negative or YYYY-MM[-DD]. Default: -31')
+@click.option('--end-date', '-ed', help='Must be negative or YYYY-MM[-DD]. Default: -1')
+@click.option('--month', '-m', help='Shortcut for -sd & -ed for a single YYYY-MM month.')
 @click.option('--where', '-w', help='WHERE conditional. Default: file.project = "project"')
 @click.option('--order', '-o', help='Field to order by. Default: download_count')
 @click.option('--all', 'all_installers', is_flag=True, help='Show downloads by all installers, not only pip.')
@@ -100,6 +102,7 @@ def pypinfo(
     days,
     start_date,
     end_date,
+    month,
     where,
     order,
     all_installers,
@@ -137,6 +140,9 @@ def pypinfo(
         order_name = order.name
         parsed_fields.insert(0, order)
 
+    if month:
+        start_date, end_date = month_ends(month)
+
     built_query = build_query(
         project,
         parsed_fields,

pypinfo/core.py

@@ -1,7 +1,8 @@
+import calendar
 import json
 import os
 import re
-from datetime import datetime
+from datetime import date, datetime
 
 from google.cloud.bigquery import Client
 from google.cloud.bigquery.job import QueryJobConfig
@@ -35,6 +36,23 @@ def normalize(name):
     return re.sub(r'[-_.]+', '-', name).lower()
 
 
+def normalize_dates(start_date, end_date):
+    """If a date is yyyy-mm, normalize as first or last yyyy-mm-dd of the month.
+    Otherwise, return unchanged.
+    """
+    try:
+        start_date, _ = month_ends(start_date)  # yyyy-mm
+    except (AttributeError, ValueError):
+        pass  # -n, yyyy-mm-dd
+
+    try:
+        _, end_date = month_ends(end_date)  # yyyy-mm
+    except (AttributeError, ValueError):
+        pass  # -n, yyyy-mm-dd
+
+    return start_date, end_date
+
+
 def create_client(creds_file=None):
     creds_file = creds_file or os.environ.get('GOOGLE_APPLICATION_CREDENTIALS')
 
@@ -59,7 +77,7 @@ def validate_date(date_text):
     except ValueError:
         pass
 
-    raise ValueError('Dates must be negative integers or YYYY-MM-DD in the past.')
+    raise ValueError('Dates must be negative integers or YYYY-MM[-DD] in the past.')
 
 
 def format_date(date, timestamp_format):
@@ -70,6 +88,15 @@ def format_date(date, timestamp_format):
     return date
 
 
+def month_ends(yyyy_mm):
+    """Helper to return start_date and end_date of a month as yyyy-mm-dd"""
+    year, month = map(int, yyyy_mm.split("-"))
+    first = date(year, month, 1)
+    number_of_days = calendar.monthrange(year, month)[1]
+    last = date(year, month, number_of_days)
+    return str(first), str(last)
+
+
 def build_query(
     project, all_fields, start_date=None, end_date=None, days=None, limit=None, where=None, order=None, pip=None
 ):
@@ -82,6 +109,7 @@ def build_query(
     if days:
         start_date = str(int(end_date) - int(days))
 
+    start_date, end_date = normalize_dates(start_date, end_date)
     validate_date(start_date)
     validate_date(end_date)
 

tests/test_core.py

@@ -118,6 +118,53 @@ def test_format_date_yyy_mm_dd():
     assert date == 'TIMESTAMP("2018-05-15 00:00:00")'
 
 
+def test_month_yyyy_mm():
+    # Act
+    first, last = core.month_ends("2019-03")
+
+    # Assert
+    assert first == "2019-03-01"
+    assert last == "2019-03-31"
+
+
+def test_month_yyyy_mm_dd():
+    # Act / Assert
+    with pytest.raises(ValueError):
+        core.month_ends("2019-03-18")
+
+
+def test_month_negative_integer():
+    # Act / Assert
+    with pytest.raises(AttributeError):
+        core.month_ends(-30)
+
+
+def test_normalize_dates_yyy_mm():
+    # Arrange
+    start_date = "2019-03"
+    end_date = "2019-03"
+
+    # Act
+    start_date, end_date = core.normalize_dates(start_date, end_date)
+
+    # Assert
+    assert start_date == "2019-03-01"
+    assert end_date == "2019-03-31"
+
+
+def test_normalize_dates_yyy_mm_dd_and_negative_integer():
+    # Arrange
+    start_date = "2019-03-18"
+    end_date = -1
+
+    # Act
+    start_date, end_date = core.normalize_dates(start_date, end_date)
+
+    # Assert
+    assert start_date == "2019-03-18"
+    assert end_date == -1
+
+
 def test_add_percentages():
     # Arrange
     rows = [