Add typespec/1 macro to transform module behaviour

lib/protobuf.ex

@@ -87,13 +87,14 @@ defmodule Protobuf do
         nil
       end
 
-      defoverridable transform_module: 0
-
       @impl unquote(__MODULE__)
       def decode(data), do: Protobuf.Decoder.decode(data, __MODULE__)
 
       @impl unquote(__MODULE__)
       def encode(struct), do: Protobuf.Encoder.encode(struct)
+
+      @on_definition {Protobuf.DSL, :on_def}
+      defoverridable transform_module: 0
     end
   end
 

lib/protobuf/dsl.ex

@@ -136,6 +136,21 @@ defmodule Protobuf.DSL do
   alias Protobuf.MessageProps
   alias Protobuf.Wire
 
+  # Registered as the @on_definition compile callback for modules that call "use Protobuf"
+  # Allow us to detect when `transform_module` is re-defined
+  def on_def(_env, :def, :transform_module, [], [], do: nil) do
+    :ok
+  end
+
+  def on_def(env, :def, :transform_module, [], [], do: module_alias_ast) do
+    Module.put_attribute(env.module, :transform_module, module_alias_ast)
+    :ok
+  end
+
+  def on_def(_, _, _, _, _, _) do
+    :ok
+  end
+
   # Registered as the @before_compile callback for modules that call "use Protobuf".
   defmacro __before_compile__(env) do
     fields = Module.get_attribute(env.module, :fields)
@@ -151,6 +166,7 @@ defmodule Protobuf.DSL do
 
     defines_t_type? = Module.defines_type?(env.module, {:t, 0})
     defines_defstruct? = Module.defines?(env.module, {:__struct__, 1})
+    transform_module_ast = Module.get_attribute(env.module, :transform_module)
 
     quote do
       @spec __message_props__() :: Protobuf.MessageProps.t()
@@ -208,7 +224,7 @@ defmodule Protobuf.DSL do
 
         # Newest version of this library generate both the t/0 type as well as the struct.
         true ->
-          unquote(def_t_typespec(msg_props, extension_props))
+          unquote(def_t_typespec(msg_props, extension_props, transform_module_ast))
           unquote(gen_defstruct(msg_props))
       end
 
@@ -226,19 +242,34 @@ defmodule Protobuf.DSL do
     end
   end
 
-  defp def_t_typespec(%MessageProps{enum?: true} = props, _extension_props) do
+  defp def_t_typespec(props, extension_props, transform_module_ast)
+       when not is_nil(transform_module_ast) do
+    default_typespec = def_t_typespec(props, extension_props, nil)
+
+    quote do
+      require unquote(transform_module_ast)
+
+      if macro_exported?(unquote(transform_module_ast), :typespec, 1) do
+        unquote(transform_module_ast).typespec(unquote(default_typespec))
+      else
+        unquote(default_typespec)
+      end
+    end
+  end
+
+  defp def_t_typespec(%MessageProps{enum?: true} = props, _extension_props, _) do
     quote do
       @type t() :: unquote(Protobuf.DSL.Typespecs.quoted_enum_typespec(props))
     end
   end
 
-  defp def_t_typespec(%MessageProps{} = props, _extension_props = nil) do
+  defp def_t_typespec(%MessageProps{} = props, _extension_props = nil, _) do
     quote do
       @type t() :: unquote(Protobuf.DSL.Typespecs.quoted_message_typespec(props))
     end
   end
 
-  defp def_t_typespec(_props, _extension_props) do
+  defp def_t_typespec(_props, _extension_props, _) do
     nil
   end
 

lib/protobuf/transform_module.ex

@@ -23,12 +23,22 @@ defmodule Protobuf.TransformModule do
       defmodule MyTransformModule do
         @behaviour Protobuf.TransformModule
 
+        defmacro typespec(_default_ast) do
+          quote do
+            @type t() :: String.t()
+          end
+        end
+
         @impl true
-        def encode(string, StringMessage) when is_binary(string), do: %StringMessage{value: string}
+        def encode(string, StringMessage) when is_binary(string), do: struct(StringMessage, value: string)
 
         @impl true
-        def decode(%StringMessage{value: string}, StringMessage), do: string
+        def decode(%{value: string}, StringMessage), do: string
       end
+
+  Notice that since the `c:typespec/1` macro was introduced, transform modules can't
+  depend on the types that they transform anymore in compile time, meaning struct
+  syntax can't be used.
   """
 
   @type value() :: term()
@@ -50,4 +60,13 @@ defmodule Protobuf.TransformModule do
   Called after a message is decoded.
   """
   @callback decode(message(), type()) :: value()
+
+  @doc """
+  Transforms the typespec for modules using this transformer.
+
+  If this callback is not present, the default typespec will be used.
+  """
+  @macrocallback typespec(default_typespec :: Macro.t()) :: Macro.t()
+
+  @optional_callbacks [typespec: 1]
 end

test/protobuf/protoc/cli_integration_test.exs

@@ -45,15 +45,15 @@ defmodule Protobuf.Protoc.CLIIntegrationTest do
       protoc!([
         "--proto_path=#{tmp_dir}",
         "--elixir_out=#{tmp_dir}",
-        "--elixir_opt=transform_module=MyTransformer",
+        "--elixir_opt=transform_module=Protobuf.TransformModule.InferFieldsFromEnum",
         "--plugin=./protoc-gen-elixir",
         proto_path
       ])
 
       assert [mod] = compile_file_and_clean_modules_on_exit("#{tmp_dir}/user.pb.ex")
       assert mod == Foo.User
 
-      assert mod.transform_module() == MyTransformer
+      assert mod.transform_module() == Protobuf.TransformModule.InferFieldsFromEnum
     end
 
     test "gen_descriptors option", %{tmp_dir: tmp_dir, proto_path: proto_path} do

test/support/test_msg.ex

@@ -249,6 +249,39 @@ defmodule TestMsg do
     field :mapsi, 3, repeated: true, map: true, type: MapFoo
   end
 
+  defmodule TransformModule do
+    @behaviour Protobuf.TransformModule
+
+    alias TestMsg.WithTransformModule
+
+    @impl true
+    defmacro typespec(default_typespec) do
+      case __CALLER__.module do
+        WithTransformModule ->
+          quote do
+            @type t() :: integer()
+          end
+
+        _ ->
+          default_typespec
+      end
+    end
+
+    # In an actual implementation, one could write the implementations of
+    # encode/2 and decode/2 in a separate module and use the structs
+    # directly.
+
+    @impl true
+    def encode(integer, WithTransformModule) when is_integer(integer) do
+      struct(WithTransformModule, field: integer)
+    end
+
+    @impl true
+    def decode(%{__struct__: WithTransformModule, field: integer}, WithTransformModule) do
+      integer
+    end
+  end
+
   defmodule WithTransformModule do
     use Protobuf, syntax: :proto3
 
@@ -277,17 +310,23 @@ defmodule TestMsg do
     field :field, 1, type: WithNewTransformModule
   end
 
-  defmodule TransformModule do
+  defmodule TransformIntegerStrings do
     @behaviour Protobuf.TransformModule
 
+    alias TestMsg.ContainsIntegerStringTransformModule
+
     @impl true
-    def encode(integer, WithTransformModule) when is_integer(integer) do
-      %WithTransformModule{field: integer}
+    def encode(
+          %{__struct__: ContainsIntegerStringTransformModule, field: str},
+          ContainsIntegerStringTransformModule
+        )
+        when is_binary(str) do
+      struct(ContainsIntegerStringTransformModule, field: String.to_integer(str))
     end
 
     @impl true
-    def decode(%WithTransformModule{field: integer}, WithTransformModule) do
-      integer
+    def decode(%{__struct__: ContainsIntegerStringTransformModule} = value, _) do
+      value
     end
   end
 
@@ -299,24 +338,6 @@ defmodule TestMsg do
     def transform_module(), do: TestMsg.TransformIntegerStrings
   end
 
-  defmodule TransformIntegerStrings do
-    @behaviour Protobuf.TransformModule
-
-    @impl true
-    def encode(
-          %ContainsIntegerStringTransformModule{field: str},
-          ContainsIntegerStringTransformModule
-        )
-        when is_binary(str) do
-      %ContainsIntegerStringTransformModule{field: String.to_integer(str)}
-    end
-
-    @impl true
-    def decode(%ContainsIntegerStringTransformModule{} = value, _) do
-      value
-    end
-  end
-
   defmodule Ext.EnumFoo do
     @moduledoc false
     use Protobuf, enum: true, syntax: :proto2