xref: /aosp_15_r20/external/tensorflow/tensorflow/python/util/function_parameter_canonicalizer.h (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1 /* Copyright 2020 The TensorFlow 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 
16 #ifndef TENSORFLOW_PYTHON_UTIL_FUNCTION_PARAMETER_CANONICALIZER_H_
17 #define TENSORFLOW_PYTHON_UTIL_FUNCTION_PARAMETER_CANONICALIZER_H_
18 
19 #include <Python.h>
20 
21 #include <vector>
22 
23 #include "absl/types/span.h"
24 #include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
25 
26 namespace tensorflow {
27 
28 // A class that Canonicalizes Python arg & kwargs parameters.
29 class FunctionParameterCanonicalizer {
30  public:
31   // `arg_names` is a list of argument names, and `defaults` is default PyObject
32   // instances for arguments. `default` is aligned to the end.
33   FunctionParameterCanonicalizer(absl::Span<const char*> arg_names,
34                                  absl::Span<PyObject*> defaults);
35 
36   // Returns the total number of arguments.
37   ABSL_MUST_USE_RESULT
GetArgSize()38   int GetArgSize() const { return interned_arg_names_.size(); }
39 
40   // Canonicalizes `args` and `kwargs` by the spec specified at construction.
41   // It's written to `result`. Returns `true` if Canonicalization was
42   // successful, and `false` otherwise. When it fails, it also sets CPython
43   // error status.
44   // This function does not update reference counter of any Python objects.
45   // `PyObject*`s in `result` are borrowed references from `args`, `kwargs`, and
46   // possibly `defaults_`, and will be only valid if `args` and `kwargs` are
47   // still alive.
48   ABSL_MUST_USE_RESULT
49   ABSL_ATTRIBUTE_HOT
50   bool Canonicalize(PyObject* args, PyObject* kwargs,
51                     absl::Span<PyObject*> result);
52 
53  private:
54   // Simple linear search of `name` in `interned_arg_names`. If found, returns
55   // the index. If not found, returns `interned_arg_names.size()`.
56   ABSL_MUST_USE_RESULT
57   ABSL_ATTRIBUTE_HOT
58   std::size_t InternedArgNameLinearSearch(PyObject* name);
59 
60   // Check if `interned_arg_names_` is unique.
61   bool AreInternedArgNamesUnique();
62 
63   // TODO(kkb): Use one `std::vector` and two `absl:Span`s instead to improve
64   // cache locality.
65   std::vector<Safe_PyObjectPtr> interned_arg_names_;
66   std::vector<Safe_PyObjectPtr> defaults_;
67   const int positional_args_size_;
68 };
69 
70 }  // namespace tensorflow
71 
72 #endif  // TENSORFLOW_PYTHON_UTIL_FUNCTION_PARAMETER_CANONICALIZER_H_
73