utils.py 8.44 KB
Newer Older
1
# Copyright 2012 - 2014 GRNET S.A. All rights reserved.
2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33
#
# Redistribution and use in source and binary forms, with or
# without modification, are permitted provided that the following
# conditions are met:
#
#   1. Redistributions of source code must retain the above
#      copyright notice, this list of conditions and the following
#      disclaimer.
#
#   2. Redistributions in binary form must reproduce the above
#      copyright notice, this list of conditions and the following
#      disclaimer in the documentation and/or other materials
#      provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
# OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
# AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
# The views and conclusions contained in the software and
# documentation are those of the authors and should not be
# interpreted as representing official policies, either expressed
# or implied, of GRNET S.A.

34
import json
35
import operator
36 37 38
import locale
import unicodedata

39 40
from datetime import datetime
from django.utils.timesince import timesince, timeuntil
41
from django.db.models.query import QuerySet
42
from django.utils.encoding import smart_unicode, smart_str
43
from snf_django.management.unicodecsv import UnicodeWriter
44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66


def smart_locale_unicode(s, **kwargs):
    """Wrapper around 'smart_unicode' using user's preferred encoding."""
    encoding = locale.getpreferredencoding()
    return smart_unicode(s, encoding=encoding, **kwargs)


def smart_locale_str(s, errors='replace', **kwargs):
    """Wrapper around 'smart_str' using user's preferred encoding."""
    encoding = locale.getpreferredencoding()
    return smart_str(s, encoding=encoding, errors=errors, **kwargs)


def safe_string(s):
    """Escape control characters from unicode and string objects."""
    if not isinstance(s, basestring):
        return s
    if isinstance(s, unicode):
        return "".join(ch.encode("unicode_escape")
                       if unicodedata.category(ch)[0] == "C" else
                       ch for ch in s)
    return s.encode("string_escape")
67 68 69 70 71


def parse_bool(value, strict=True):
    """Convert a string to boolen value.

72
    If strict is True, then ValueError will be raised, if the string can not be
73 74 75
    converted to boolean. Otherwise the string will be returned as is.

    """
76 77 78
    if isinstance(value, bool):
        return value

79 80 81 82 83 84
    if value.lower() in ("yes", "true", "t", "1"):
        return True
    if value.lower() in ("no", "false", "f", "0"):
        return False

    if strict:
85
        raise ValueError("Cannot convert '%s' to boolean value" % value)
86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104
    else:
        return value


def format_bool(b):
    """Convert a boolean value to YES or NO."""
    return "YES" if b else "NO"


def format_date(d):
    if not d:
        return ""

    if d < datetime.now():
        return timesince(d) + " ago"
    else:
        return "in " + timeuntil(d)


105 106 107 108 109 110 111 112 113 114
def filter_results(results, filters):
    if isinstance(results, QuerySet):
        return filter_queryset_results(results, filters)
    elif isinstance(results, list):
        return filter_object_results(results, filters)
    else:
        raise ValueError("Invalid type for results argument: %s", results)


def parse_queryset_filters(filters):
115 116 117 118 119 120 121 122 123
    """Parse a string into lookup parameters for QuerySet.filter(**kwargs).

    This functions converts a string of comma-separated key 'cond' val triples
    to two dictionaries, containing lookup parameters to be used for filter
    and exclude functions of QuerySet.

    e.g. filter_by="foo>=2, baz!=4" -> ({"foo__gte": "2"}, {"baz": "4"})

    """
124 125 126 127 128 129 130 131 132
    OP_MAP = [
        (">=", "__gte"),
        ("=>", "__gte"),
        (">",  "__gt"),
        ("<=", "__lte"),
        ("=<", "__lte"),
        ("<", "__lt"),
        ("=", ""),
        ]
133 134 135

    filter_dict = {}
    exclude_dict = {}
136 137 138
    for filter_str in filters.split(","):
        if "!=" in filter_str:
            key, val = filter_str.split("!=")
139
            exclude_dict[key] = parse_bool(val, strict=False)
140
            continue
141
        for op, new_op in OP_MAP:
142 143
            if op in filter_str:
                key, val = filter_str.split(op)
144
                filter_dict[key + new_op] = parse_bool(val, strict=False)
145 146 147
                break
        else:
            raise ValueError("Unknown filter expression: %s" % filter_str)
148 149 150 151

    return (filter_dict, exclude_dict)


152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191
def filter_queryset_results(results, filters):
    filter_dict, exclude_dict = parse_queryset_filters(filters)
    return results.exclude(**exclude_dict).filter(**filter_dict)


def parse_object_filters(filters):
    OP_MAP = [
        (">=", operator.ge),
        ("=>", operator.ge),
        (">",  operator.gt),
        ("<=", operator.le),
        ("=<", operator.le),
        ("<", operator.lt),
        ("!=", operator.ne),
        ("=", operator.eq),
    ]
    filters = []
    for filter_str in filters.split(","):
        for op, op_func in OP_MAP:
            if op in filter_str:
                key, val = filter_str.split(op)
                filters.append((key.strip(), op_func, val.strip()))
                break
        else:
            raise ValueError("Unknown filter expression: %s" % filter_str)
    return filters


def filter_object_results(results, filters):
    results = list(results)
    if results is []:
        return results
    zero_result = results[0]
    for key, op_func, val in parse_object_filters(filters):
        val_type = type(getattr(zero_result, key))
        results = filter(lambda x: op_func(getattr(x, key), val_type(val)),
                         results)
    return results


192
def pprint_table(out, table, headers=None, output_format='pretty',
193
                 separator=None, vertical=False, title=None):
194 195 196 197 198 199 200 201 202 203
    """Print a pretty, aligned string representation of table.

    Works by finding out the max width of each column and padding to data
    to this value.
    """

    assert(isinstance(table, (list, tuple))), "Invalid table type"
    if headers:
        assert(isinstance(headers, (list, tuple))), "Invalid headers type"

204 205
    sep = separator if separator else "  "

206
    if headers:
207 208
        headers = map(smart_unicode, headers)
    table = [map(smart_unicode, row) for row in table]
209 210

    if output_format == "json":
211
        assert(headers is not None), "json output format requires headers"
212 213 214 215
        table = [dict(zip(headers, row)) for row in table]
        out.write(json.dumps(table, indent=4))
        out.write("\n")
    elif output_format == "csv":
216 217
        enc = locale.getpreferredencoding()
        cw = UnicodeWriter(out, encoding=enc)
218
        if headers:
219 220
            table.insert(0, headers)
        cw.writerows(table)
221
    elif output_format == "pretty":
222 223 224 225 226 227
        if vertical:
            assert(len(table) == 1)
            row = table[0]
            max_key = max(map(len, headers))
            for row in table:
                for (k, v) in zip(headers, row):
228
                    k = k.ljust(max_key)
229 230 231 232 233 234 235
                    out.write("%s: %s\n" % (k, v))
        else:
            # Find out the max width of each column
            columns = [headers] + table if headers else table
            widths = [max(map(len, col)) for col in zip(*(columns))]

            t_length = sum(widths) + len(sep) * (len(widths) - 1)
236
            if title is not None:
237
                t_length = max(t_length, len(title))
238
                out.write("-" * t_length + "\n")
239
                out.write(title.center(t_length) + "\n")
240
                out.write("-" * t_length + "\n")
241 242
            if headers:
                # pretty print the headers
243
                line = sep.join(v.rjust(w)
244 245 246 247 248 249
                                for v, w in zip(headers, widths))
                out.write(line + "\n")
                out.write("-" * t_length + "\n")

            # print the rest table
            for row in table:
250
                line = sep.join(v.rjust(w) for v, w in zip(row, widths))
251
                out.write(line + "\n")
252 253
    else:
        raise ValueError("Unknown output format '%s'" % output_format)