Django Model自定义IF、IFNULL函数

1. 新建一个myfunction.py文件

# coding: utf-8

from django.db.models import Expression, Q, F


class If(Expression):
    template = 'IF(%(condition)s, %(then)s, %(_else)s)'

    def __init__(self, condition=None, then=None, _else=None, output_field=None, **lookups):
        if lookups and condition is None:
            condition, lookups = Q(**lookups), None
        if condition is None or not isinstance(condition, Q) or lookups:
            raise TypeError("__init__() takes either a Q object or lookups as keyword arguments")

        self.condition = condition
        self.then = self._parse_expressions(then)[0]
        self._else = self._parse_expressions(_else)[0]

        if output_field is not None:
            self.output_field = output_field

    def get_source_expressions(self):
        return [self.condition, self.then, self._else]

    def set_source_expressions(self, exprs):
        self.condition, self.then, self._else = exprs

    def get_source_fields(self):
        return [self.then._output_field_or_none, self._else._output_field_or_none]

    def as_sql(self, compiler, connection, template=None, **extra_context):
        connection.ops.check_expression_support(self)
        template_params = extra_context

        sql_params = []
        condition_sql, condition_params = compiler.compile(self.condition)
        template_params['condition'] = condition_sql
        sql_params.extend(condition_params)

        then_sql, then_params = compiler.compile(self.then)
        template_params['then'] = then_sql
        sql_params.extend(then_params)

        else_sql, else_params = compiler.compile(self._else)
        template_params['_else'] = else_sql
        sql_params.extend(else_params)

        template = template or self.template
        return template % template_params, sql_params

    def __str__(self):
        return "IF(%r, %r, %r)" % (self.condition, self.then, self._else)

    def __repr__(self):
        return "<%s: %s>" % (self.__class__.__name__, self)


class IfNull(Expression):
    template = 'IFNULL(%(condition)s, %(then)s)'

    def __init__(self, condition, then, output_field=None):
        if isinstance(condition, str):
            condition = F(condition)

        self.condition = condition
        self.then = self._parse_expressions(then)[0]

        if output_field is not None:
            self.output_field = output_field

    def get_source_expressions(self):
        return [self.condition, self.then]

    def set_source_expressions(self, exprs):
        self.condition, self.then = exprs

    def get_source_fields(self):
        return [self.then._output_field_or_none]

    def as_sql(self, compiler, connection, template=None, **extra_context):
        connection.ops.check_expression_support(self)
        template_params = extra_context

        sql_params = []
        condition_sql, condition_params = compiler.compile(self.condition)
        template_params['condition'] = condition_sql
        sql_params.extend(condition_params)

        then_sql, then_params = compiler.compile(self.then)
        template_params['then'] = then_sql
        sql_params.extend(then_params)

        template = template or self.template
        return template % template_params, sql_params

    def __str__(self):
        return "IFNULL(%r, %r)" % (self.condition, self.then)

    def __repr__(self):
        return "<%s: %s>" % (self.__class__.__name__, self)

2. 使用

from function import If, IfNull
from django.db.models import Q, Value

# if的用法，sql为： if(field1 is null, field1, field2) as new_field
YouTable.objects.values(new_field = If(field1__isnull=True, 'field2', 'field1'))

# if的用法，sql为： if(field1 is null and field2 > 0 and , field3, 0) as new_field
YouTable.objects.values(new_field = If(Q(field1__isnull=True) & Q(field2__gt=0), 'field3', Value(0))

# ifnull的用法，sql为： ifnull(field1, field2) as new_field
YouTable.objects.values(new_field = IfNull('field1', 'field2'))

# ifnull的用法，sql为： ifnull(field1, 0) as new_field
YouTable.objects.values(new_field = IfNull(F('field1'), Value(0)))