Repository: altair-viz/altair
Pull request #866 (2018-05-17 21:34:13)

ENH: parse timeUnits in encoding shorthands
Addresses #849 (cc/ @kanitw)

This changes the ``parse_shorthand`` utility so that it parses timeUnits as well as aggregates.

This allows compact expressions of time-binned encodings; e.g.
```python
import altair as alt
from vega_datasets import data

alt.Chart(data.seattle_weather.url).mark_rect().encode(
    x='date(date):O',
    y='month(date):O',
    color='max(temp_max):Q'
)
```
![visualization 2](https://user-images.githubusercontent.com/781659/40204974-3388afbc-59df-11e8-84b8-fa54063e6d7e.png)

The old way of specifying this (which does still work) is:
```python
alt.Chart(data.seattle_weather.url).mark_rect().encode(
    x=alt.X('date:O', timeUnit='date'),
    y=alt.Y('date:O', timeUnit='month'),
    color='max(temp_max):Q'
)
```

As an additional benefit, the new approach also makes more clear that the typecode refers to the aggregated data rather than to the input data, which was a confusion that came up several times for participants in the PyCon tutorial.


Patch:
diff --git a/altair/utils/core.py b/altair/utils/core.py
index eea9d118..b659ed14 100644
--- a/altair/utils/core.py
+++ b/altair/utils/core.py
@@ -1,12 +1,13 @@
 """
 Utility routines
 """
-import re
-import warnings
 import collections
 from copy import deepcopy
+import itertools
+import re
 import sys
 import traceback
+import warnings
 
 import six
 import pandas as pd
@@ -28,6 +29,31 @@ TYPECODE_MAP = {'ordinal': 'O',
 INV_TYPECODE_MAP = {v: k for k, v in TYPECODE_MAP.items()}
 
 
+# aggregates from vega-lite version 2.4.3
+AGGREGATES = ['argmax', 'argmin', 'average', 'count', 'distinct', 'max',
+              'mean', 'median', 'min', 'missing', 'q1', 'q3', 'ci0', 'ci1',
+              'stderr', 'stdev', 'stdevp', 'sum', 'valid', 'values',
+              'variance', 'variancep']
+
+# timeUnits from vega-lite version 2.4.3
+TIMEUNITS = ["utcyear", "utcquarter", "utcmonth", "utcday", "utcdate",
+             "utchours", "utcminutes", "utcseconds", "utcmilliseconds",
+             "utcyearquarter", "utcyearquartermonth", "utcyearmonth",
+             "utcyearmonthdate", "utcyearmonthdatehours",
+             "utcyearmonthdatehoursminutes",
+             "utcyearmonthdatehoursminutesseconds",
+             "utcquartermonth", "utcmonthdate", "utchoursminutes",
+             "utchoursminutesseconds", "utcminutesseconds",
+             "utcsecondsmilliseconds",
+             "year", "quarter", "month", "day", "date", "hours", "minutes",
+             "seconds", "milliseconds", "yearquarter", "yearquartermonth",
+             "yearmonth", "yearmonthdate", "yearmonthdatehours",
+             "yearmonthdatehoursminutes",
+             "yearmonthdatehoursminutesseconds", "quartermonth", "monthdate",
+             "hoursminutes", "hoursminutesseconds", "minutesseconds",
+             "secondsmilliseconds"]
+
+
 def infer_vegalite_type(data):
     """
     From an array-like input, infer the correct vega typecode
@@ -107,9 +133,9 @@ def sanitize_dataframe(df):
     return df
 
 
-def _parse_shorthand(shorthand):
-    """
-    Parse the shorthand expression for aggregation, field, and type.
+def parse_shorthand(shorthand, data=None, parse_aggregates=True,
+                    parse_timeunits=True, parse_types=True):
+    """General tool to parse shorthand values
 
     These are of the form:
 
@@ -118,115 +144,116 @@ def _parse_shorthand(shorthand):
     - "average(col_name)"
     - "average(col_name):O"
 
+    Optionally, a dataframe may be supplied, from which the type
+    will be inferred if not specified in the shorthand.
+
     Parameters
     ----------
-    shorthand: str
-        Shorthand string
+    shorthand : dict or string
+        The shorthand representation to be parsed
+    data : DataFrame, optional
+        If specified and of type DataFrame, then use these values to infer the
+        column type if not provided by the shorthand.
+    parse_aggregates : boolean
+        If True (default), then parse aggregate functions within the shorthand.
+    parse_timeunits : boolean
+        If True (default), then parse timeUnits from within the shorthand
+    parse_types : boolean
+        If True (default), then parse typecodes within the shorthand
 
     Returns
     -------
-    D : dict
-        Dictionary containing the field, aggregate, and typecode
-    """
-    if not shorthand:
-        return {}
+    attrs : dict
+        a dictionary of attributes extracted from the shorthand
 
-    # List taken from vega-lite v2 AggregateOp
-    valid_aggregates = ["argmax", "argmin", "average", "count", "distinct",
-                        "max", "mean", "median", "min", "missing", "q1", "q3",
-                        "ci0", "ci1", "stderr", "stdev", "stdevp", "sum",
-                        "valid", "values", "variance", "variancep"]
-    valid_typecodes = list(TYPECODE_MAP) + list(INV_TYPECODE_MAP)
+    Examples
+    --------
+    >>> data = pd.DataFrame({'foo': ['A', 'B', 'A', 'B'],
+    ...                      'bar': [1, 2, 3, 4]})
 
-    # build regular expressions
-    units = dict(field='(?P<field>.*)',
-                 type='(?P<type>{0})'.format('|'.join(valid_typecodes)),
-                 count='(?P<aggregate>count)',
-                 aggregate='(?P<aggregate>{0})'.format('|'.join(valid_aggregates)))
-    patterns = [r'{count}\(\)',
-                r'{count}\(\):{type}',
-                r'{aggregate}\({field}\):{type}',
-                r'{aggregate}\({field}\)',
-                r'{field}:{type}',
-                r'{field}']
-    regexps = (re.compile('\A' + p.format(**units) + '\Z', re.DOTALL)
-               for p in patterns)
+    >>> parse_shorthand('name') == {'field': 'name'}
+    True
 
-    # find matches depending on valid fields passed
-    match = next(exp.match(shorthand).groupdict() for exp in regexps
-                 if exp.match(shorthand))
+    >> parse_shorthand('name:Q') == {'field': 'name', 'type': 'quantitative'}
+    True
 
-    # Handle short form of the type expression
-    type_ = match.get('type', None)
-    if type_:
-        match['type'] = INV_TYPECODE_MAP.get(type_, type_)
+    >>> parse_shorthand('average(col)') == {'aggregate': 'average', 'field': 'col'}
+    True
 
-    # counts are quantitative by default
-    if match == {'aggregate': 'count'}:
-        match['type'] = 'quantitative'
+    >>> parse_shorthand('foo:O') == {'field': 'foo', 'type': 'ordinal'}
+    True
 
-    return match
+    >>> parse_shorthand('min(foo):Q') == {'aggregate': 'min', 'field': 'foo', 'type': 'quantitative'}
+    True
 
+    >>> parse_shorthand('month(col)') == {'field': 'col', 'timeUnit': 'month', 'type': 'temporal'}
+    True
 
-def parse_shorthand(shorthand, data=None):
-    """Parse the shorthand expression for aggregation, field, and type.
+    >>> parse_shorthand('year(col):O') == {'field': 'col', 'timeUnit': 'year', 'type': 'ordinal'}
+    True
 
-    These are of the form:
+    >>> parse_shorthand('foo', data) == {'field': 'foo', 'type': 'nominal'}
+    True
 
-    - "col_name"
-    - "col_name:O"
-    - "average(col_name)"
-    - "average(col_name):O"
+    >>> parse_shorthand('bar', data) == {'field': 'bar', 'type': 'quantitative'}
+    True
 
-    Optionally, a dataframe may be supplied, from which the type
-    will be inferred if not specified in the shorthand.
+    >>> parse_shorthand('bar:O', data) == {'field': 'bar', 'type': 'ordinal'}
+    True
 
-    Parameters
-    ----------
-    shorthand: str
-        Shorthand string of the form "agg(col):typ"
-    data : pd.DataFrame (optional)
-        Dataframe from which to infer types
+    >>> parse_shorthand('sum(bar)', data) == {'aggregate': 'sum', 'field': 'bar', 'type': 'quantitative'}
+    True
 
-    Returns
-    -------
-    D : dict
-        Dictionary which always contains a 'field' key, and additionally
-        contains an 'aggregate' and 'type' key depending on the input.
+    >>> parse_shorthand('count()', data) == {'aggregate': 'count', 'type': 'quantitative'}
+    True
+    """
+    if not shorthand:
+        return {}
 
-    Examples
-    --------
-    >>> data = pd.DataFrame({'foo': ['A', 'B', 'A', 'B'],
-    ...                      'bar': [1, 2, 3, 4]})
+    valid_typecodes = list(TYPECODE_MAP) + list(INV_TYPECODE_MAP)
+
+    units = dict(field='(?P<field>.*)',
+                 type='(?P<type>{0})'.format('|'.join(valid_typecodes)),
+                 count='(?P<aggregate>count)',
+                 aggregate='(?P<aggregate>{0})'.format('|'.join(AGGREGATES)),
+                 timeUnit='(?P<timeUnit>{0})'.format('|'.join(TIMEUNITS)))
+
+    patterns = []
 
-    >>> parse_shorthand('name')
-    {'field': 'name'}
+    if parse_aggregates:
+        patterns.extend([r'{count}\(\)',
+                         r'{aggregate}\({field}\)'])
+    if parse_timeunits:
+        patterns.extend([r'{timeUnit}\({field}\)'])
 
-    >>> parse_shorthand('average(col)')  # doctest: +SKIP
-    {'aggregate': 'average', 'field': 'col'}
+    patterns.extend([r'{field}'])
 
-    >>> parse_shorthand('foo:O')  # doctest: +SKIP
-    {'field': 'foo', 'type': 'ordinal'}
+    if parse_types:
+        patterns = list(itertools.chain(*((p + ':{type}', p) for p in patterns)))
 
-    >>> parse_shorthand('min(foo):Q')  # doctest: +SKIP
-    {'aggregate': 'min', 'field': 'foo', 'type': 'quantitative'}
+    regexps = (re.compile('\A' + p.format(**units) + '\Z', re.DOTALL)
+               for p in patterns)
 
-    >>> parse_shorthand('foo', data)  # doctest: +SKIP
-    {'field': 'foo', 'type': 'nominal'}
+    # find matches depending on valid fields passed
+    if isinstance(shorthand, dict):
+        attrs = shorthand
+    else:
+        attrs = next(exp.match(shorthand).groupdict() for exp in regexps
+                     if exp.match(shorthand))
 
-    >>> parse_shorthand('bar', data)  # doctest: +SKIP
-    {'field': 'bar', 'type': 'quantitative'}
+    # Handle short form of the type expression
+    if 'type' in attrs:
+        attrs['type'] = INV_TYPECODE_MAP.get(attrs['type'], attrs['type'])
 
-    >>> parse_shorthand('bar:O', data)  # doctest: +SKIP
-    {'field': 'bar', 'type': 'ordinal'}
+    # counts are quantitative by default
+    if attrs == {'aggregate': 'count'}:
+        attrs['type'] = 'quantitative'
 
-    >>> parse_shorthand('sum(bar)', data)  # doctest: +SKIP
-    {'aggregate': 'sum', 'field': 'bar', 'type': 'quantitative'}
+    # times are temporal by default
+    if 'timeUnit' in attrs and 'type' not in attrs:
+        attrs['type'] = 'temporal'
 
-    >>> parse_shorthand('count()', data)  # doctest: +SKIP
-    {'aggregate': 'count', 'type': 'quantitative'}
-    """
-    attrs = _parse_shorthand(shorthand)
+    # if data is specified and type is not, infer type from data
     if isinstance(data, pd.DataFrame) and 'type' not in attrs:
         if 'field' in attrs and attrs['field'] in data.columns:
             attrs['type'] = infer_vegalite_type(data[attrs['field']])
diff --git a/altair/vegalite/v2/api.py b/altair/vegalite/v2/api.py
index e5e92f52..6695b159 100644
--- a/altair/vegalite/v2/api.py
+++ b/altair/vegalite/v2/api.py
@@ -805,23 +805,66 @@ class TopLevelMixin(mixins.ConfigMethodMixin):
             The data field to apply time unit.
         timeUnit : TimeUnit
             The timeUnit.
+        **kwargs
+            transforms can also be passed by keyword argument; see Examples
 
         Returns
         -------
         self : Chart object
             returns chart to allow for chaining
 
+        Examples
+        --------
+        >>> import altair as alt
+        >>> from altair import datum, expr
+
+        >>> chart = alt.Chart().transform_timeunit(month='month(date)')
+        >>> chart.transform[0]
+        TimeUnitTransform({
+          as: 'month',
+          field: 'date',
+          timeUnit: 'month'
+        })
+
+        It's also possible to pass the ``TimeUnitTransform`` arguments directly;
+        this is most useful in cases where the desired field name is not a
+        valid python identifier:
+
+        >>> kwds = {'as': 'month', 'timeUnit': 'month', 'field': 'The Month'}
+        >>> chart = alt.Chart().transform_timeunit(**kwds)
+        >>> chart.transform[0]
+        TimeUnitTransform({
+          as: 'month',
+          field: 'The Month',
+          timeUnit: 'month'
+        })
+
+        As the first form is easier to write and understand, that is the
+        recommended method.
+
         See Also
         --------
         alt.TimeUnitTransform : underlying transform object
         """
-        if as_ is not Undefined:
+        if as_ is Undefined:
+            as_ = kwargs.pop('as', Undefined)
+        else:
             if 'as' in kwargs:
                 raise ValueError("transform_timeunit: both 'as_' and 'as' passed as arguments.")
-            kwargs['as'] = as_
-        kwargs['field'] = field
-        kwargs['timeUnit'] = timeUnit
-        return self._add_transform(core.TimeUnitTransform(**kwargs))
+        if as_ is not Undefined:
+            dct = {'as': as_, 'timeUnit': timeUnit, 'field': field}
+            self = self._add_transform(core.TimeUnitTransform(**dct))
+        for as_, shorthand in kwargs.items():
+            dct = utils.parse_shorthand(shorthand,
+                                        parse_timeunits=True,
+                                        parse_aggregates=False,
+                                        parse_types=False)
+            dct.pop('type', None)
+            dct['as'] = as_
+            if 'timeUnit' not in dct:
+                raise ValueError("'{0}' must include a valid timeUnit".format(shorthand))
+            self = self._add_transform(core.TimeUnitTransform(**dct))
+        return self
 
     @utils.use_signature(core.WindowTransform)
     def transform_window(self, *args, **kwargs):
diff --git a/doc/user_guide/transform.rst b/doc/user_guide/transform.rst
index 063f0941..f14f8c19 100644
--- a/doc/user_guide/transform.rst
+++ b/doc/user_guide/transform.rst
@@ -587,12 +587,12 @@ measurements in Seattle during the year 2010:
 
 The plot is too busy due to the amount of data points squeezed into the short
 time; we can make it a bit cleaner by discretizing it, for example, by month
-(``timeUnit="month"``) and plotting only the mean monthly temperature:
+and plotting only the mean monthly temperature:
 
 .. altair-plot::
 
     alt.Chart(temps).mark_line().encode(
-        alt.X('date:T', timeUnit='month'),
+        x='month(date):T',
         y='mean(temp):Q'
     )
 
@@ -603,7 +603,7 @@ This can be useful when plotting a bar chart or other discrete chart type:
 .. altair-plot::
 
     alt.Chart(temps).mark_bar().encode(
-        alt.X('date:O', timeUnit='month'),
+        x='month(date):O',
         y='mean(temp):Q'
     )
 
@@ -614,8 +614,8 @@ to give a profile of Seattle temperatures through the year:
 .. altair-plot::
 
     alt.Chart(temps).mark_rect().encode(
-        alt.X('date:O', timeUnit='date', axis=alt.Axis(title='day')),
-        alt.Y('date:O', timeUnit='month', axis=alt.Axis(title='month')),
+        alt.X('date(date):O', title='day'),
+        alt.Y('month(date):O', title='month'),
         color='max(temp):Q'
     ).properties(
         title="2010 Daily High Temperatures in Seattle (F)"
@@ -634,7 +634,7 @@ method. For example:
         alt.X('month:T', axis=alt.Axis(format='%b')),
         y='mean(temp):Q'
     ).transform_timeunit(
-        "month", field="date", timeUnit="month"
+        month='month(date)'
     )
 
 Notice that because the ``timeUnit`` is not part of the encoding channel here,


Test patch:
diff --git a/altair/utils/tests/test_core.py b/altair/utils/tests/test_core.py
index 07f9115b..3a6029fc 100644
--- a/altair/utils/tests/test_core.py
+++ b/altair/utils/tests/test_core.py
@@ -1,5 +1,6 @@
 import pandas as pd
 
+import altair as alt
 from .. import parse_shorthand, update_nested
 
 
@@ -42,6 +43,13 @@ def test_parse_shorthand():
     check('count()', aggregate='count', type='quantitative')
     check('count():O', aggregate='count', type='ordinal')
 
+    # time units:
+    check('month(x)', field='x', timeUnit='month', type='temporal')
+    check('year(foo):O', field='foo', timeUnit='year', type='ordinal')
+    check('date(date):quantitative',
+          field='date', timeUnit='date', type='quantitative')
+    check('yearmonthdate(field)', field='field', timeUnit='yearmonthdate', type='temporal')
+
 
 def test_parse_shorthand_with_data():
     def check(s, data, **kwargs):
@@ -55,7 +63,30 @@ def test_parse_shorthand_with_data():
     check('y', data, field='y', type='nominal')
     check('z', data, field='z', type='temporal')
     check('count(x)', data, field='x', aggregate='count', type='quantitative')
-    check('mean(*)', data, field='*', aggregate='mean')
+    check('count()', data, aggregate='count', type='quantitative')
+    check('month(z)', data, timeUnit='month', field='z', type='temporal')
+
+
+def test_parse_shorthand_all_aggregates():
+    aggregates = alt.Root._schema['definitions']['AggregateOp']['enum']
+    for aggregate in aggregates:
+        shorthand = "{aggregate}(field):Q".format(aggregate=aggregate)
+        assert parse_shorthand(shorthand) == {'aggregate': aggregate,
+                                              'field': 'field',
+                                              'type': 'quantitative'}
+
+
+def test_parse_shorthand_all_timeunits():
+    timeUnits = []
+    for loc in ['Local', 'Utc']:
+        for typ in ['Single', 'Multi']:
+            defn = loc + typ + 'TimeUnit'
+            timeUnits.extend(alt.Root._schema['definitions'][defn]['enum'])
+    for timeUnit in timeUnits:
+        shorthand = "{timeUnit}(field):Q".format(timeUnit=timeUnit)
+        assert parse_shorthand(shorthand) == {'timeUnit': timeUnit,
+                                              'field': 'field',
+                                              'type': 'quantitative'}
 
 
 def test_update_nested():