#!/usr/bin/python

#

# Copyright (C) 2009 Google Inc.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

#      http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

"""Converts Python data into data for Google Visualization API clients.

This library can be used to create a google.visualization.DataTable usable by

visualizations built on the Google Visualization API. Output formats are raw

JSON, JSON response, and JavaScript.

See http://code.google.com/apis/visualization/ for documentation on the

Google Visualization API.

"""

__author__ = "Amit Weinstein, Misha Seltzer"

import cgi

import datetime

import types

class DataTableException(Exception):

  """The general exception object thrown by DataTable."""

  pass

class DataTable(object):

  """Wraps the data to convert to a Google Visualization API DataTable.

  Create this object, populate it with data, then call one of the ToJS...

  methods to return a string representation of the data in the format described.

  You can clear all data from the object to reuse it, but you cannot clear

  individual cells, rows, or columns. You also cannot modify the table schema

  specified in the class constructor.

  You can add new data one or more rows at a time. All data added to an

  instantiated DataTable must conform to the schema passed in to __init__().

  You can reorder the columns in the output table, and also specify row sorting

  order by column. The default column order is according to the original

  table_description parameter. Default row sort order is ascending, by column

  1 values. For a dictionary, we sort the keys for order.

  The data and the table_description are closely tied, as described here:

  The table schema is defined in the class constructor's table_description

  parameter. The user defines each column using a tuple of

  (id[, type[, label[, custom_properties]]]). The default value for type is

  string, label is the same as ID if not specified, and custom properties is

  an empty dictionary if not specified.

  table_description is a dictionary or list, containing one or more column

  descriptor tuples, nested dictionaries, and lists. Each dictionary key, list

  element, or dictionary element must eventually be defined as

  a column description tuple. Here's an example of a dictionary where the key

  is a tuple, and the value is a list of two tuples:

    {('a', 'number'): [('b', 'number'), ('c', 'string')]}

  This flexibility in data entry enables you to build and manipulate your data

  in a Python structure that makes sense for your program.

  Add data to the table using the same nested design as the table's

  table_description, replacing column descriptor tuples with cell data, and

  each row is an element in the top level collection. This will be a bit

  clearer after you look at the following examples showing the

  table_description, matching data, and the resulting table:

  Columns as list of tuples [col1, col2, col3]

    table_description: [('a', 'number'), ('b', 'string')]

    AppendData( [[1, 'z'], [2, 'w'], [4, 'o'], [5, 'k']] )

    Table:

    a  b   <--- these are column ids/labels

    1  z

    2  w

    4  o

    5  k

  Dictionary of columns, where key is a column, and value is a list of

  columns  {col1: [col2, col3]}

    table_description: {('a', 'number'): [('b', 'number'), ('c', 'string')]}

    AppendData( data: {1: [2, 'z'], 3: [4, 'w']}

    Table:

    a  b  c

    1  2  z

    3  4  w

  Dictionary where key is a column, and the value is itself a dictionary of

  columns {col1: {col2, col3}}

    table_description: {('a', 'number'): {'b': 'number', 'c': 'string'}}

    AppendData( data: {1: {'b': 2, 'c': 'z'}, 3: {'b': 4, 'c': 'w'}}

    Table:

    a  b  c

    1  2  z

    3  4  w

  """

  def __init__(self, table_description, data=None, custom_properties=None):

    """Initialize the data table from a table schema and (optionally) data.

    See the class documentation for more information on table schema and data

    values.

    Args:

      table_description: A table schema, following one of the formats described

                         in TableDescriptionParser(). Schemas describe the

                         column names, data types, and labels. See

                         TableDescriptionParser() for acceptable formats.

      data: Optional. If given, fills the table with the given data. The data

            structure must be consistent with schema in table_description. See

            the class documentation for more information on acceptable data. You

            can add data later by calling AppendData().

      custom_properties: Optional. A dictionary from string to string that

                         goes into the table's custom properties. This can be

                         later changed by changing self.custom_properties.

    Raises:

      DataTableException: Raised if the data and the description did not match,

                          or did not use the supported formats.

    """

    self.__columns = self.TableDescriptionParser(table_description)

    self.__data = []

    self.custom_properties = {}

    if custom_properties is not None:

      self.custom_properties = custom_properties

    if data:

      self.LoadData(data)

  @staticmethod

  def _EscapeValueForCsv(v):

    """Escapes the value for use in a CSV file.

    Puts the string in double-quotes, and escapes any inner double-quotes by

    doubling them.

    Args:

      v: The value to escape.

    Returns:

      The escaped values.

    """

    return '"%s"' % v.replace('"', '""')

  @staticmethod

  def _EscapeValue(v):

    """Puts the string in quotes, and escapes any inner quotes and slashes."""

    if isinstance(v, unicode):

      # Here we use repr as in the usual case, but on unicode strings, it

      # also escapes the unicode characters (which we want to leave as is).

      # So, after repr() we decode using raw-unicode-escape, which decodes

      # only the unicode characters, and leaves all the rest (", ', \n and

      # more) escaped.

      # We don't take the first character, because repr adds a u in the

      # beginning of the string (usual repr output for unicode is u'...').

      return repr(v).decode("raw-unicode-escape")[1:]

    # Here we use python built-in escaping mechanism for string using repr.

    return repr(str(v))

  @staticmethod

  def _EscapeCustomProperties(custom_properties):

    """Escapes the custom properties dictionary."""

    l = []

    for key, value in custom_properties.iteritems():

      l.append("%s:%s" % (DataTable._EscapeValue(key),

                          DataTable._EscapeValue(value)))

    return "{%s}" % ",".join(l)

  @staticmethod

  def SingleValueToJS(value, value_type, escape_func=None):

    """Translates a single value and type into a JS value.

    Internal helper method.

    Args:

      value: The value which should be converted

      value_type: One of "string", "number", "boolean", "date", "datetime" or

                  "timeofday".

      escape_func: The function to use for escaping strings.

    Returns:

      The proper JS format (as string) of the given value according to the

      given value_type. For None, we simply return "null".

      If a tuple is given, it should be in one of the following forms:

        - (value, formatted value)

        - (value, formatted value, custom properties)

      where the formatted value is a string, and custom properties is a

      dictionary of the custom properties for this cell.

      To specify custom properties without specifying formatted value, one can

      pass None as the formatted value.

      One can also have a null-valued cell with formatted value and/or custom

      properties by specifying None for the value.

      This method ignores the custom properties except for checking that it is a

      dictionary. The custom properties are handled in the ToJSon and ToJSCode

      methods.

      The real type of the given value is not strictly checked. For example,

      any type can be used for string - as we simply take its str( ) and for

      boolean value we just check "if value".

      Examples:

        SingleValueToJS(None, "boolean") returns "null"

        SingleValueToJS(False, "boolean") returns "false"

        SingleValueToJS((5, "5$"), "number") returns ("5", "'5$'")

        SingleValueToJS((None, "5$"), "number") returns ("null", "'5$'")

    Raises:

      DataTableException: The value and type did not match in a not-recoverable

                          way, for example given value 'abc' for type 'number'.

    """

    if escape_func is None:

      escape_func = DataTable._EscapeValue

    if isinstance(value, tuple):

      # In case of a tuple, we run the same function on the value itself and

      # add the formatted value.

      if (len(value) not in [2, 3] or

          (len(value) == 3 and not isinstance(value[2], dict))):

        raise DataTableException("Wrong format for value and formatting - %s." %

                                 str(value))

      if not isinstance(value[1], types.StringTypes + (types.NoneType,)):

        raise DataTableException("Formatted value is not string, given %s." %

                                 type(value[1]))

      js_value = DataTable.SingleValueToJS(value[0], value_type)

      if value[1] is None:

        return (js_value, None)

      return (js_value, escape_func(value[1]))

    # The standard case - no formatting.

    t_value = type(value)

    if value is None:

      return "null"

    if value_type == "boolean":

      if value:

        return "true"

      return "false"

    elif value_type == "number":

      if isinstance(value, (int, long, float)):

        return str(value)

      raise DataTableException("Wrong type %s when expected number" % t_value)

    elif value_type == "string":

      if isinstance(value, tuple):

        raise DataTableException("Tuple is not allowed as string value.")

      return escape_func(value)

    elif value_type == "date":

      if not isinstance(value, (datetime.date, datetime.datetime)):

        raise DataTableException("Wrong type %s when expected date" % t_value)

        # We need to shift the month by 1 to match JS Date format

      return "new Date(%d,%d,%d)" % (value.year, value.month - 1, value.day)

    elif value_type == "timeofday":

      if not isinstance(value, (datetime.time, datetime.datetime)):

        raise DataTableException("Wrong type %s when expected time" % t_value)

      return "[%d,%d,%d]" % (value.hour, value.minute, value.second)

    elif value_type == "datetime":

      if not isinstance(value, datetime.datetime):

        raise DataTableException("Wrong type %s when expected datetime" %

                                 t_value)

      return "new Date(%d,%d,%d,%d,%d,%d)" % (value.year,

                                              value.month - 1,  # To match JS

                                              value.day,

                                              value.hour,

                                              value.minute,

                                              value.second)

    # If we got here, it means the given value_type was not one of the

    # supported types.

    raise DataTableException("Unsupported type %s" % value_type)

  @staticmethod

  def ColumnTypeParser(description):

    """Parses a single column description. Internal helper method.

    Args:

      description: a column description in the possible formats:

       'id'

       ('id',)

       ('id', 'type')

       ('id', 'type', 'label')

       ('id', 'type', 'label', {'custom_prop1': 'custom_val1'})

    Returns:

      Dictionary with the following keys: id, label, type, and

      custom_properties where:

        - If label not given, it equals the id.

        - If type not given, string is used by default.

        - If custom properties are not given, an empty dictionary is used by

          default.

    Raises:

      DataTableException: The column description did not match the RE.

    """

    if not description:

      raise DataTableException("Description error: empty description given")

    if not isinstance(description, (types.StringTypes, tuple)):

      raise DataTableException("Description error: expected either string or "

                               "tuple, got %s." % type(description))

    if isinstance(description, types.StringTypes):

      description = (description,)

    # According to the tuple's length, we fill the keys

    # We verify everything is of type string

    for elem in description[:3]:

      if not isinstance(elem, types.StringTypes):

        raise DataTableException("Description error: expected tuple of "

                                 "strings, current element of type %s." %

                                 type(elem))

    desc_dict = {"id": description[0],

                 "label": description[0],

                 "type": "string",

                 "custom_properties": {}}

    if len(description) > 1:

      desc_dict["type"] = description[1].lower()

      if len(description) > 2:

        desc_dict["label"] = description[2]

        if len(description) > 3:

          if not isinstance(description[3], dict):

            raise DataTableException("Description error: expected custom "

                                     "properties of type dict, current element "

                                     "of type %s." % type(description[3]))

          desc_dict["custom_properties"] = description[3]

          if len(description) > 4:

            raise DataTableException("Description error: tuple of length > 4")

    return desc_dict

  @staticmethod

  def TableDescriptionParser(table_description, depth=0):

    """Parses the table_description object for internal use.

    Parses the user-submitted table description into an internal format used

    by the Python DataTable class. Returns the flat list of parsed columns.

    Args:

      table_description: A description of the table which should comply

                         with one of the formats described below.

      depth: Optional. The depth of the first level in the current description.

             Used by recursive calls to this function.

    Returns:

      List of columns, where each column represented by a dictionary with the

      keys: id, label, type, depth, container which means the following:

      - id: the id of the column

      - name: The name of the column

      - type: The datatype of the elements in this column. Allowed types are

              described in ColumnTypeParser().

      - depth: The depth of this column in the table description

      - container: 'dict', 'iter' or 'scalar' for parsing the format easily.

      - custom_properties: The custom properties for this column.

      The returned description is flattened regardless of how it was given.

    Raises:

      DataTableException: Error in a column description or in the description

                          structure.

    Examples:

      A column description can be of the following forms:

       'id'

       ('id',)

       ('id', 'type')

       ('id', 'type', 'label')

       ('id', 'type', 'label', {'custom_prop1': 'custom_val1'})

       or as a dictionary:

       'id': 'type'

       'id': ('type',)

       'id': ('type', 'label')

       'id': ('type', 'label', {'custom_prop1': 'custom_val1'})

      If the type is not specified, we treat it as string.

      If no specific label is given, the label is simply the id.

      If no custom properties are given, we use an empty dictionary.

      input: [('a', 'date'), ('b', 'timeofday', 'b', {'foo': 'bar'})]

      output: [{'id': 'a', 'label': 'a', 'type': 'date',

                'depth': 0, 'container': 'iter', 'custom_properties': {}},

               {'id': 'b', 'label': 'b', 'type': 'timeofday',

                'depth': 0, 'container': 'iter',

                'custom_properties': {'foo': 'bar'}}]

      input: {'a': [('b', 'number'), ('c', 'string', 'column c')]}

      output: [{'id': 'a', 'label': 'a', 'type': 'string',

                'depth': 0, 'container': 'dict', 'custom_properties': {}},

               {'id': 'b', 'label': 'b', 'type': 'number',

                'depth': 1, 'container': 'iter', 'custom_properties': {}},

               {'id': 'c', 'label': 'column c', 'type': 'string',

                'depth': 1, 'container': 'iter', 'custom_properties': {}}]

      input:  {('a', 'number', 'column a'): { 'b': 'number', 'c': 'string'}}

      output: [{'id': 'a', 'label': 'column a', 'type': 'number',

                'depth': 0, 'container': 'dict', 'custom_properties': {}},

               {'id': 'b', 'label': 'b', 'type': 'number',

                'depth': 1, 'container': 'dict', 'custom_properties': {}},

               {'id': 'c', 'label': 'c', 'type': 'string',

                'depth': 1, 'container': 'dict', 'custom_properties': {}}]

      input: { ('w', 'string', 'word'): ('c', 'number', 'count') }

      output: [{'id': 'w', 'label': 'word', 'type': 'string',

                'depth': 0, 'container': 'dict', 'custom_properties': {}},

               {'id': 'c', 'label': 'count', 'type': 'number',

                'depth': 1, 'container': 'scalar', 'custom_properties': {}}]

    """

    # For the recursion step, we check for a scalar object (string or tuple)

    if isinstance(table_description, (types.StringTypes, tuple)):

      parsed_col = DataTable.ColumnTypeParser(table_description)

      parsed_col["depth"] = depth

      parsed_col["container"] = "scalar"

      return [parsed_col]

    # Since it is not scalar, table_description must be iterable.

    if not hasattr(table_description, "__iter__"):

      raise DataTableException("Expected an iterable object, got %s" %

                               type(table_description))

    if not isinstance(table_description, dict):

      # We expects a non-dictionary iterable item.

      columns = []

      for desc in table_description:

        parsed_col = DataTable.ColumnTypeParser(desc)

        parsed_col["depth"] = depth

        parsed_col["container"] = "iter"

        columns.append(parsed_col)

      if not columns:

        raise DataTableException("Description iterable objects should not"

                                 " be empty.")

      return columns

    # The other case is a dictionary

    if not table_description:

      raise DataTableException("Empty dictionaries are not allowed inside"

                               " description")

    # The number of keys in the dictionary separates between the two cases of

    # more levels below or this is the most inner dictionary.

    if len(table_description) != 1:

      # This is the most inner dictionary. Parsing types.

      columns = []

      # We sort the items, equivalent to sort the keys since they are unique

      for key, value in sorted(table_description.items()):

        # We parse the column type as (key, type) or (key, type, label) using

        # ColumnTypeParser.

        if isinstance(value, tuple):

          parsed_col = DataTable.ColumnTypeParser((key,) + value)

        else:

          parsed_col = DataTable.ColumnTypeParser((key, value))

        parsed_col["depth"] = depth

        parsed_col["container"] = "dict"

        columns.append(parsed_col)

      return columns

    # This is an outer dictionary, must have at most one key.

    parsed_col = DataTable.ColumnTypeParser(table_description.keys()[0])

    parsed_col["depth"] = depth

    parsed_col["container"] = "dict"

    return ([parsed_col] +

            DataTable.TableDescriptionParser(table_description.values()[0],

                                             depth=depth + 1))

  @property

  def columns(self):

    """Returns the parsed table description."""

    return self.__columns

  def NumberOfRows(self):

    """Returns the number of rows in the current data stored in the table."""

    return len(self.__data)

  def SetRowsCustomProperties(self, rows, custom_properties):

    """Sets the custom properties for given row(s).

    Can accept a single row or an iterable of rows.

    Sets the given custom properties for all specified rows.

    Args:

      rows: The row, or rows, to set the custom properties for.

      custom_properties: A string to string dictionary of custom properties to

      set for all rows.

    """

    if not hasattr(rows, "__iter__"):

      rows = [rows]

    for row in rows:

      self.__data[row] = (self.__data[row][0], custom_properties)

  def LoadData(self, data, custom_properties=None):

    """Loads new rows to the data table, clearing existing rows.

    May also set the custom_properties for the added rows. The given custom

    properties dictionary specifies the dictionary that will be used for *all*

    given rows.

    Args:

      data: The rows that the table will contain.

      custom_properties: A dictionary of string to string to set as the custom

                         properties for all rows.

    """

    self.__data = []

    self.AppendData(data, custom_properties)

  def AppendData(self, data, custom_properties=None):

    """Appends new data to the table.

    Data is appended in rows. Data must comply with

    the table schema passed in to __init__(). See SingleValueToJS() for a list

    of acceptable data types. See the class documentation for more information

    and examples of schema and data values.

    Args:

      data: The row to add to the table. The data must conform to the table

            description format.

      custom_properties: A dictionary of string to string, representing the

                         custom properties to add to all the rows.

    Raises:

      DataTableException: The data structure does not match the description.

    """

    # If the maximal depth is 0, we simply iterate over the data table

    # lines and insert them using _InnerAppendData. Otherwise, we simply

    # let the _InnerAppendData handle all the levels.

    if not self.__columns[-1]["depth"]:

      for row in data:

        self._InnerAppendData(({}, custom_properties), row, 0)

    else:

      self._InnerAppendData(({}, custom_properties), data, 0)