xref: /aosp_15_r20/external/bazel-skylib/lib/partial.bzl (revision bcb5dc7965af6ee42bf2f21341a2ec00233a8c8a)
1# Copyright 2018 The Bazel Authors. All rights reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#    http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15"""Starlark module for working with partial function objects.
16
17Partial function objects allow some parameters are bound before the call.
18
19Similar to https://docs.python.org/3/library/functools.html#functools.partial.
20"""
21
22# create instance singletons to avoid unnecessary allocations
23_a_dict_type = type({})
24_a_tuple_type = type(())
25_a_struct_type = type(struct())
26
27def _call(partial, *args, **kwargs):
28    """Calls a partial created using `make`.
29
30    Args:
31      partial: The partial to be called.
32      *args: Additional positional arguments to be appended to the ones given to
33             make.
34      **kwargs: Additional keyword arguments to augment and override the ones
35                given to make.
36
37    Returns:
38      Whatever the function in the partial returns.
39    """
40    function_args = partial.args + args
41    function_kwargs = dict(partial.kwargs)
42    function_kwargs.update(kwargs)
43    return partial.function(*function_args, **function_kwargs)
44
45def _make(func, *args, **kwargs):
46    """Creates a partial that can be called using `call`.
47
48    A partial can have args assigned to it at the make site, and can have args
49    passed to it at the call sites.
50
51    A partial 'function' can be defined with positional args and kwargs:
52
53      # function with no args
54      ```
55      def function1():
56        ...
57      ```
58
59      # function with 2 args
60      ```
61      def function2(arg1, arg2):
62        ...
63      ```
64
65      # function with 2 args and keyword args
66      ```
67      def function3(arg1, arg2, x, y):
68        ...
69      ```
70
71    The positional args passed to the function are the args passed into make
72    followed by any additional positional args given to call. The below example
73    illustrates a function with two positional arguments where one is supplied by
74    make and the other by call:
75
76      # function demonstrating 1 arg at make site, and 1 arg at call site
77      ```
78      def _foo(make_arg1, func_arg1):
79        print(make_arg1 + " " + func_arg1 + "!")
80      ```
81
82    For example:
83
84      ```
85      hi_func = partial.make(_foo, "Hello")
86      bye_func = partial.make(_foo, "Goodbye")
87      partial.call(hi_func, "Jennifer")
88      partial.call(hi_func, "Dave")
89      partial.call(bye_func, "Jennifer")
90      partial.call(bye_func, "Dave")
91      ```
92
93    prints:
94
95      ```
96      "Hello, Jennifer!"
97      "Hello, Dave!"
98      "Goodbye, Jennifer!"
99      "Goodbye, Dave!"
100      ```
101
102    The keyword args given to the function are the kwargs passed into make
103    unioned with the keyword args given to call. In case of a conflict, the
104    keyword args given to call take precedence. This allows you to set a default
105    value for keyword arguments and override it at the call site.
106
107    Example with a make site arg, a call site arg, a make site kwarg and a
108    call site kwarg:
109
110      ```
111      def _foo(make_arg1, call_arg1, make_location, call_location):
112        print(make_arg1 + " is from " + make_location + " and " +
113              call_arg1 + " is from " + call_location + "!")
114
115      func = partial.make(_foo, "Ben", make_location="Hollywood")
116      partial.call(func, "Jennifer", call_location="Denver")
117      ```
118
119    Prints "Ben is from Hollywood and Jennifer is from Denver!".
120
121      ```
122      partial.call(func, "Jennifer", make_location="LA", call_location="Denver")
123      ```
124
125    Prints "Ben is from LA and Jennifer is from Denver!".
126
127    Note that keyword args may not overlap with positional args, regardless of
128    whether they are given during the make or call step. For instance, you can't
129    do:
130
131    ```
132    def foo(x):
133      pass
134
135    func = partial.make(foo, 1)
136    partial.call(func, x=2)
137    ```
138
139    Args:
140      func: The function to be called.
141      *args: Positional arguments to be passed to function.
142      **kwargs: Keyword arguments to be passed to function. Note that these can
143                be overridden at the call sites.
144
145    Returns:
146      A new `partial` that can be called using `call`
147    """
148    return struct(function = func, args = args, kwargs = kwargs)
149
150def _is_instance(v):
151    """Returns True if v is a partial created using `make`.
152
153    Args:
154      v: The value to check.
155
156    Returns:
157      True if v was created by `make`, False otherwise.
158    """
159
160    # Note that in bazel 3.7.0 and earlier, type(v.function) is the same
161    # as the type of a function even if v.function is a rule. But we
162    # cannot rely on this in later bazels due to breaking change
163    # https://github.com/bazelbuild/bazel/commit/e379ece1908aafc852f9227175dd3283312b4b82
164    #
165    # Since this check is heuristic anyway, we simply check for the
166    # presence of a "function" attribute without checking its type.
167    return type(v) == _a_struct_type and \
168           hasattr(v, "function") and \
169           hasattr(v, "args") and type(v.args) == _a_tuple_type and \
170           hasattr(v, "kwargs") and type(v.kwargs) == _a_dict_type
171
172partial = struct(
173    make = _make,
174    call = _call,
175    is_instance = _is_instance,
176)
177