Support reserved fields in enums (#205)

Author: Drvi

src/codegen/metadata_methods.jl

@@ -10,6 +10,11 @@ function maybe_generate_reserved_fields_method(io, t::MessageType)
     println(io, "PB.reserved_fields(::Type{", safename(t), "}) = (names = ", string(t.reserved_names), ", numbers = Union{Int,UnitRange{Int}}[", join(t.reserved_nums, ", "), "])")
 end
 
+function maybe_generate_reserved_fields_method(io, t::EnumType)
+    isempty(t.reserved_names) && isempty(t.reserved_nums) && return
+    println(io, "PB.reserved_fields(::Type{", safename(t), ".T}) = (names = ", string(t.reserved_names), ", numbers = Union{Int,UnitRange{Int}}[", join(t.reserved_nums, ", "), "])")
+end
+
 function maybe_generate_extendable_field_numbers_method(io, t::MessageType)
     n = length(t.extensions)
     n == 0 && return

src/codegen/toplevel_definitions.jl

@@ -120,6 +120,7 @@ function codegen(io, t::EnumType, ::Context)
     end
     println(io)
     maybe_generate_deprecation(io, t)
+    maybe_generate_reserved_fields_method(io, t)
 end
 
 function codegen(io, t::ServiceType, ::Context)

src/parsing/proto_types.jl

@@ -117,6 +117,8 @@ struct EnumType <: AbstractProtoType
     element_names::Vector{Symbol}
     element_values::Vector{Int}
     options::Dict{String,Union{String,Dict{String}}}
+    reserved_nums::Vector{Union{Int,UnitRange{Int}}}
+    reserved_names::Vector{String}
     field_options::Dict{String,Union{String,Dict{String}}}
 end
 
@@ -268,6 +270,8 @@ function parse_enum_type(ps::ParserState, name_prefix="")
     name = unsafe_name(ps)
 
     options = Dict{String,Union{String,Dict{String}}}()
+    reserved_nums = Vector{Union{Int,UnitRange{Int}}}()
+    reserved_names = Vector{String}()
     field_options = Dict{String,Union{String,Dict{String}}}()
     element_names = Symbol[]
     element_values = Int[]
@@ -286,6 +290,8 @@ function parse_enum_type(ps::ParserState, name_prefix="")
                 parse_field_options!(ps, get!(field_options, element_name, Dict{String,Union{String,Dict{String}}}()))
             end
             expectnext(ps, Tokens.SEMICOLON)
+        elseif accept(ps, Tokens.RESERVED)
+            _parse_reserved_statement!(ps, reserved_nums, reserved_names)
         else
             expectnext(ps, Tokens.RBRACE)
             break
@@ -300,7 +306,7 @@ function parse_enum_type(ps::ParserState, name_prefix="")
         error("In proto3, enums' first element must map to zero, $name has `$(first(element_names)) = $(first(element_values))` as first element.")
     end
     # TODO: validate field_numbers
-    return EnumType(_dot_join(name_prefix, name), element_names, element_values, options, field_options)
+    return EnumType(_dot_join(name_prefix, name), element_names, element_values, options, reserved_nums, reserved_names, field_options)
 end
 
 # We consumed EXTEND

test/test_codegen.jl

@@ -5,6 +5,7 @@ using ProtoBuf.CodeGenerators: resolve_inter_package_references!, get_all_transi
 using ProtoBuf.CodeGenerators: CodeGenerators
 using ProtoBuf.Parsers: parse_proto_file, ParserState, Parsers
 using ProtoBuf.Lexers: Lexer
+using EnumX
 using Test
 
 strify(f, args...) = (io = IOBuffer(); f(io, args...); String(take!(io)))
@@ -634,15 +635,18 @@ end
 
     @testset "Metadata methods" begin
         @testset "metadata_methods have generic fallback" begin
-            s, p, ctx = translate_simple_proto("message A { }")
+            s, p, ctx = translate_simple_proto("message A { } enum Foo { }")
             @test strify(CodeGenerators.maybe_generate_reserved_fields_method, p.definitions["A"]) == ""
+            @test strify(CodeGenerators.maybe_generate_reserved_fields_method, p.definitions["Foo"]) == ""
             @test strify(CodeGenerators.maybe_generate_extendable_field_numbers_method, p.definitions["A"]) == ""
             @test strify(CodeGenerators.maybe_generate_default_values_method, p.definitions["A"], ctx) == ""
             @test strify(CodeGenerators.maybe_generate_oneof_field_types_method, p.definitions["A"], ctx) == ""
             @test strify(CodeGenerators.maybe_generate_field_numbers_method, p.definitions["A"]) == ""
 
             struct A end
+            @enumx Foo
             @test reserved_fields(A) == (names = String[], numbers = Union{Int,UnitRange{Int}}[])
+            @test reserved_fields(Foo.T) == (names = String[], numbers = Union{Int,UnitRange{Int}}[])
             @test extendable_field_numbers(A) == Union{Int,UnitRange{Int}}[]
             @test default_values(A) == (;)
             @test oneof_field_types(A) == (;)
@@ -658,6 +662,16 @@ end
             @test strify(CodeGenerators.maybe_generate_field_numbers_method,            p.definitions["A"])      == "PB.field_numbers(::Type{A}) = (;a = 1, s = 3)\n"
             @test strify(CodeGenerators.maybe_generate_kwarg_constructor_method,        p.definitions["A"], ctx) == "A(;a = nothing, o = nothing) = A(a, o)\n"
         end
+
+        @testset "reserved fields are available for enums" begin
+            s, p, ctx = translate_simple_proto("""
+            enum Foo {
+                reserved 2, 15, 9 to 11, 40 to max;
+                reserved "FOO", "BAR";
+              }
+            """)
+            @test strify(CodeGenerators.maybe_generate_reserved_fields_method, p.definitions["Foo"]) =="PB.reserved_fields(::Type{Foo.T}) = (names = [\"FOO\", \"BAR\"], numbers = Union{Int,UnitRange{Int}}[2, 15, 9:11, 40:536870911])\n"
+        end
     end
 
     @testset "Imports within a leaf module which name-clashes with top module" begin