1# -*- coding: utf-8 -*-

2

3"""

4***************************************************************************

5 qgsfunction.py

6 ---------------------

7 Date : May 2018

8 Copyright : (C) 2018 by Denis Rouzaud

9 Email : denis@opengis.ch

10***************************************************************************

11* *

12* This program is free software; you can redistribute it and/or modify *

13* it under the terms of the GNU General Public License as published by *

14* the Free Software Foundation; either version 2 of the License, or *

15* (at your option) any later version. *

16* *

17***************************************************************************

18"""

19

20

21 import inspect

22 import string

23 from builtins import str

24 from qgis.PyQt.QtCore import QCoreApplication

25 from qgis._core import QgsExpressionFunction, QgsExpression, QgsMessageLog, QgsFeatureRequest, Qgis

26

27

28 def register_function(function, arg_count, group, usesgeometry=False,

29 referenced_columns=[QgsFeatureRequest.ALL_ATTRIBUTES], handlesnull=False, **kwargs):

30 """

31 Register a Python function to be used as a expression function.

32

33 Functions should take (values, feature, parent) as args:

34

35 Example:

36 def myfunc(values, feature, parent):

37 pass

38

39 They can also shortcut naming feature and parent args by using *args

40 if they are not needed in the function.

41

42 Example:

43 def myfunc(values, *args):

44 pass

45

46 Functions should return a value compatible with QVariant

47

48 Eval errors can be raised using parent.setEvalErrorString("Error message")

49

50 :param function:

51 :param arg_count:

52 :param group:

53 :param usesgeometry:

54 :param handlesnull: Needs to be set to True if this function does not always return NULL if any parameter is NULL. Default False.

55 :return:

56 """

57

58 class QgsPyExpressionFunction(QgsExpressionFunction):

59

60 def __init__(self, func, name, args, group, helptext='', usesGeometry=True,

61 referencedColumns=QgsFeatureRequest.ALL_ATTRIBUTES, expandargs=False, handlesNull=False):

62 QgsExpressionFunction.__init__(self, name, args, group, helptext)

63 self.function = func

64 self.expandargs = expandargs

65 self.uses_geometry = usesGeometry

66 self.referenced_columns = referencedColumns

67 self.handles_null = handlesNull

68

69 def func(self, values, context, parent, node):

70 feature = None

71 if context:

72 feature = context.feature()

73

74 try:

75 if self.expandargs:

76 values.append(feature)

77 values.append(parent)

78 if inspect.getfullargspec(self.function).args[-1] == 'context':

79 values.append(context)

80 return self.function(*values)

81 else:

82 if inspect.getfullargspec(self.function).args[-1] == 'context':

83 self.function(values, feature, parent, context)

84 return self.function(values, feature, parent)

85 except Exception as ex:

86 parent.setEvalErrorString(str(ex))

87 return None

88

89 def usesGeometry(self, node):

90 return self.uses_geometry

91

92 def referencedColumns(self, node):

93 return self.referenced_columns

94

95 def handlesNull(self):

96 return self.handles_null

97

98 helptemplate = string.Template("""<h3>$name function</h3><br>$doc""")

99 name = kwargs.get('name', function.__name__)

100 helptext = kwargs.get('helpText') or function.__doc__ or ''

101 helptext = helptext.strip()

102 expandargs = False

103

104 if arg_count == "auto":

105 # Work out the number of args we need.

106 # Number of function args - 2. The last two args are always feature, parent.

107 args = inspect.getfullargspec(function).args

108 number = len(args)

109 arg_count = number - 2

110 if args[-1] == 'context':

111 arg_count -= 1

112 expandargs = True

113

114 register = kwargs.get('register', True)

115 if register and QgsExpression.isFunctionName(name):

116 if not QgsExpression.unregisterFunction(name):

117 msgtitle = QCoreApplication.translate("UserExpressions", "User expressions")

118 msg = QCoreApplication.translate("UserExpressions",

119 "The user expression {0} already exists and could not be unregistered.").format(

120 name)

121 QgsMessageLog.logMessage(msg + "\n", msgtitle, Qgis.Warning)

122 return None

123

124 function.__name__ = name

125 helptext = helptemplate.safe_substitute(name=name, doc=helptext)

126 f = QgsPyExpressionFunction(function, name, arg_count, group, helptext, usesgeometry, referenced_columns,

127 expandargs, handlesnull)

128

129 # This doesn't really make any sense here but does when used from a decorator context

130 # so it can stay.

131 if register:

132 QgsExpression.registerFunction(f)

133 return f

134

135

136 def qgsfunction(args='auto', group='custom', **kwargs):

137 """

138 Decorator function used to define a user expression function.

139

140 :param args: Number of parameters, set to 'auto' to accept a variable length of parameters.

141 :param group: The expression group to which this expression should be added.

142 :param \**kwargs:

143 See below

144

145 :Keyword Arguments:

146 * *referenced_columns* (``list``) --

147 An array of field names on which this expression works. Can be set to ``[QgsFeatureRequest.ALL_ATTRIBUTES]``. By default empty.

148 * *usesgeometry* (``bool``) --

149 Defines if this expression requires the geometry. By default False.

150 * *handlesnull* (``bool``) --

151 Defines if this expression has custom handling for NULL values. If False, the result will always be NULL as soon as any parameter is NULL. False by default.

152

153 Example:

154 @qgsfunction(2, 'test'):

155 def add(values, feature, parent):

156 pass

157

158 Will create and register a function in QgsExpression called 'add' in the

159 'test' group that takes two arguments.

160

161 or not using feature and parent:

162

163 Example:

164 @qgsfunction(2, 'test'):

165 def add(values, *args):

166 pass

167 """

168

169 def wrapper(func):

170 return register_function(func, args, group, **kwargs)

171

172 return wrapper