THRIFT-5859: Generate a values iterator for enum types for go
Client: go
The generated code looks like this:
type Foo int64
const (
Foo_One Foo = 1
Foo_Two Foo = 2
)
var knownFooValues = []Foo{
Foo_One,
Foo_Two,
}
func FooValues() iter.Seq[Foo] {
return func(yield func(Foo) bool) {
for _, v := range knownFooValues {
if !yield(v) {
return
}
}
}
}
diff --git a/compiler/cpp/src/thrift/generate/t_go_generator.cc b/compiler/cpp/src/thrift/generate/t_go_generator.cc
index 9e1865f..8d05892 100644
--- a/compiler/cpp/src/thrift/generate/t_go_generator.cc
+++ b/compiler/cpp/src/thrift/generate/t_go_generator.cc
@@ -706,6 +706,7 @@
}
system_packages.push_back("errors");
system_packages.push_back("fmt");
+ system_packages.push_back("iter");
system_packages.push_back("log/slog");
system_packages.push_back("time");
// For the thrift import, always do rename import to make sure it's called thrift.
@@ -731,6 +732,7 @@
"var _ = context.Background\n"
"var _ = errors.New\n"
"var _ = fmt.Printf\n"
+ "var _ = iter.Pull[int]\n"
"var _ = slog.Log\n"
"var _ = time.Now\n"
"var _ = thrift.ZERO\n"
@@ -784,11 +786,13 @@
* @param tenum The enumeration
*/
void t_go_generator::generate_enum(t_enum* tenum) {
- std::ostringstream to_string_mapping, from_string_mapping;
+ std::ostringstream to_string_mapping, from_string_mapping, known_values_mapping;
std::string tenum_name(publicize(tenum->get_name()));
generate_go_docstring(f_types_, tenum);
generate_deprecation_comment(f_types_, tenum->annotations_);
- f_types_ << "type " << tenum_name << " int64" << '\n' << "const (" << '\n';
+ f_types_ << "type " << tenum_name << " int64" << '\n' << '\n' << "const (" << '\n';
+
+ known_values_mapping << indent() << "var known" << tenum_name << "Values" << " = []" << tenum_name << "{" << '\n';
to_string_mapping << indent() << "func (p " << tenum_name << ") String() string {" << '\n';
indent_up();
@@ -815,6 +819,7 @@
generate_deprecation_comment(f_types_, (*c_iter)->annotations_);
f_types_ << indent() << tenum_name << "_" << iter_name << ' ' << tenum_name << " = "
<< value << '\n';
+ known_values_mapping << indent() << tenum_name << "_" << iter_name << "," << '\n';
// Dictionaries to/from string names of enums
to_string_mapping << indent() << "case " << tenum_name << "_" << iter_name << ": return \""
<< iter_std_name << "\"" << '\n';
@@ -840,8 +845,29 @@
indent_down();
from_string_mapping << indent() << "}" << '\n';
- f_types_ << ")" << '\n' << '\n' << to_string_mapping.str() << '\n' << from_string_mapping.str()
- << '\n' << '\n';
+ known_values_mapping << indent() << "}" << '\n' << '\n';
+ known_values_mapping << indent() << "func " << tenum_name << "Values() iter.Seq[" << tenum_name << "] {" << '\n';
+ indent_up();
+ known_values_mapping << indent() << "return func(yield func(" << tenum_name << ") bool) {" << '\n';
+ indent_up();
+ known_values_mapping << indent() << "for _, v := range known" << tenum_name << "Values {" << '\n';
+ indent_up();
+ known_values_mapping << indent() << "if !yield(v) {" << '\n';
+ indent_up();
+ known_values_mapping << indent() << "return" << '\n';
+ indent_down();
+ known_values_mapping << indent() << "}" << '\n';
+ indent_down();
+ known_values_mapping << indent() << "}" << '\n';
+ indent_down();
+ known_values_mapping << indent() << "}" << '\n';
+ indent_down();
+ known_values_mapping << indent() << "}" << '\n';
+
+ f_types_ << ")" << '\n' << '\n'
+ << known_values_mapping.str() << '\n'
+ << to_string_mapping.str() << '\n'
+ << from_string_mapping.str() << '\n' << '\n';
// Generate a convenience function that converts an instance of an enum
// (which may be a constant) into a pointer to an instance of that enum
diff --git a/lib/go/test/tests/enum_values_test.go b/lib/go/test/tests/enum_values_test.go
new file mode 100644
index 0000000..84bd265
--- /dev/null
+++ b/lib/go/test/tests/enum_values_test.go
@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package tests
+
+import (
+ "slices"
+ "testing"
+
+ "github.com/apache/thrift/lib/go/test/gopath/src/constoptionalfielda"
+)
+
+func TestEnumValues(t *testing.T) {
+ want := []constoptionalfielda.Foo{
+ constoptionalfielda.Foo_One,
+ constoptionalfielda.Foo_Two,
+ }
+ got := slices.Collect(constoptionalfielda.FooValues())
+ t.Logf("FooValues = %#v", got)
+ if len(got) != len(want) {
+ t.Fatalf("Want %d values in FooValues(), got %+v", len(want), got)
+ }
+ for i, v := range want {
+ if got[i] != v {
+ t.Errorf("FooValues()[%d] got %v(%d) want %v(%d)", i, got[i], got[i], v, v)
+ }
+ }
+}