Thrift generator for Starlark constants

Summary:
Enhance Thrift generator for Starlark adding support for constants. Only a subset of types are supported:
- primitive types (string, integers, floating point numbers)
- lists
- maps
- enums

Structs (out of scope) and sets (Starlark doesn't support sets) are not supported.

Reviewed By: vitaut

Differential Revision: D52435866

fbshipit-source-id: c168a8e55de2e01db6bf064c7db6d88e1e76af9d

thrift/compiler/generate/t_starlark_generator.cc

@@ -31,15 +31,69 @@ class t_starlark_generator : public t_mstch_generator {
   std::string template_prefix() const override { return "starlark"; }
 
   void generate_program() override;
+
+ private:
+  void set_mstch_factories();
+};
+
+class mstch_starlark_type : public mstch_type {
+ public:
+  mstch_starlark_type(
+      const t_type* t, mstch_context& ctx, mstch_element_position pos)
+      : mstch_type(t, ctx, pos) {
+    register_methods(
+        this,
+        {
+            {"type:starlark_supported?",
+             &mstch_starlark_type::is_starlark_supported},
+        });
+  }
+
+  mstch::node is_starlark_supported() { return is_supported_type(type_); }
+
+ private:
+  bool is_supported_type(const t_type* type) {
+    return is_primitive(type->get_true_type()) ||
+        is_supported_collection(type->get_true_type());
+  }
+
+  bool is_primitive(const t_type* type) {
+    return type->is_byte() || type->is_i16() || type->is_i32() ||
+        type->is_i64() || type->is_string() || type->is_bool() ||
+        type->is_float() || type->is_double();
+  }
+
+  bool is_supported_collection(const t_type* type) {
+    if (type->is_enum()) {
+      return true;
+    }
+    if (type->is_list()) {
+      return is_supported_type(
+          dynamic_cast<const t_list*>(type)->get_elem_type());
+    }
+    if (type->is_map()) {
+      const t_map* map = dynamic_cast<const t_map*>(type);
+      return is_supported_type(map->get_key_type()) &&
+          is_supported_type(map->get_val_type());
+    }
+    return false;
+  }
 };
 
 void t_starlark_generator::generate_program() {
   out_dir_base_ = "gen-star";
+
+  set_mstch_factories();
   const auto* program = get_program();
   auto mstch_program = mstch_context_.program_factory->make_mstch_object(
       program, mstch_context_);
+
   render_to_file(
-      std::move(mstch_program), "enums.star", program->name() + ".star");
+      std::move(mstch_program), "definitions.star", program->name() + ".star");
+}
+
+void t_starlark_generator::set_mstch_factories() {
+  mstch_context_.add<mstch_starlark_type>();
 }
 
 THRIFT_REGISTER_GENERATOR(starlark, "Starlark", "Starlark generator");

thrift/compiler/generate/templates/starlark/consts.mustache

@@ -0,0 +1,12 @@
+{{#program:constants}}{{!
+  }}{{#constant:type}}{{!
+    }}{{#type:starlark_supported?}}
+{{constant:name}} = {{!
+      }}{{#constant:value}}{{!
+        }}{{> consts/value}}{{!
+      }}{{/constant:value}}
+      {{!newline after constant}}
+
+    {{/type:starlark_supported?}}{{!
+  }}{{/constant:type}}{{!
+}}{{/program:constants}}

thrift/compiler/generate/templates/starlark/consts/value.mustache

@@ -0,0 +1,54 @@
+{{#type:byte?}}{{value:integer_value}}{{/type:byte?}}{{!
+}}{{#type:i16?}}{{value:integer_value}}{{/type:i16?}}{{!
+}}{{#type:i32?}}{{value:integer_value}}{{/type:i32?}}{{!
+}}{{#type:i64?}}{{value:integer_value}}{{/type:i64?}}{{!
+}}{{#type:string?}}"{{value:string_value}}"{{/type:string?}}{{!
+}}{{#type:bool?}}{{!
+  }}{{#value:nonzero?}}True{{/value:nonzero?}}{{!
+  }}{{^value:nonzero?}}False{{/value:nonzero?}}{{!
+}}{{/type:bool?}}{{!
+}}{{#type:float?}}{{!
+  }}{{#value:double?}}{{value:double_value}}{{/value:double?}}{{!
+  }}{{#value:integer?}}{{value:integer_value}}.0{{/value:integer?}}{{!
+}}{{/type:float?}}{{!
+}}{{#type:double?}}{{!
+  }}{{#value:double?}}{{value:double_value}}{{/value:double?}}{{!
+  }}{{#value:integer?}}{{value:integer_value}}.0{{/value:integer?}}{{!
+}}{{/type:double?}}{{!
+
+  List constants
+}}{{#type:list?}}[
+{{#value:list_elements}}
+  {{#type:list_elem_type}}{{!
+    }}{{> consts/value}},{{!
+  }}{{/type:list_elem_type}}
+{{/value:list_elements}}
+]{{/type:list?}}{{!
+
+  Map constants
+}}{{#type:map?}}{
+{{#value:map_elements}}
+  {{#element:key}}{{!
+    }}{{#type:key_type}}{{!
+      }}{{> consts/value}}{{!
+    }}{{/type:key_type}}{{!
+  }}{{/element:key}}: {{!
+  }}{{#element:value}}{{!
+    }}{{#type:value_type}}{{!
+      }}{{> consts/value}}{{!
+    }}{{/type:value_type}}{{!
+  }}{{/element:value}},
+{{/value:map_elements}}
+}{{/type:map?}}{{!
+
+  Enum constants
+}}{{#type:enum?}}{{!
+  }}{{#type:enum}}{{!
+    }}{{#value:enum_value?}}{{!
+      }}{{#value:enum_value}}{{!
+        Known enum value with a name.
+        }}{{enum:name}}.{{enum_value:name}}{{!
+      }}{{/value:enum_value}}{{!
+    }}{{/value:enum_value?}}{{!
+  }}{{/type:enum}}{{!
+}}{{/type:enum?}}

thrift/compiler/generate/templates/starlark/definitions.star.mustache

@@ -0,0 +1,2 @@
+{{> enums}}
+{{> consts}}

thrift/compiler/test/fixtures/constants/cmd

@@ -13,3 +13,4 @@ mstch_swift src/module.thrift
 json src/module.thrift
 mstch_python src/module.thrift
 mstch_python_capi src/module.thrift
+starlark src/module.thrift