feat: add type parameters

Author: fahchen

.formatter.exs

@@ -1,4 +1,4 @@
-locals_without_parens = [field: 2, field: 3, plugin: 1, plugin: 2]
+locals_without_parens = [parameter: 1, field: 2, field: 3, plugin: 1, plugin: 2]
 
 [
   inputs: [

README.md

@@ -182,6 +182,34 @@ defmodule MyModule do
 end
 ```
 
+You can also add type parameters:
+```elixir
+defmodule User do
+  use TypedStruct
+
+  typedstruct do
+    parameter :state
+
+    field :state, state, enforce: true
+    field :name, String.t()
+  end
+end
+```
+
+this equals to:
+```elixir
+defmodule User do
+  @enforce_keys [:state]
+  defstruct state: nil,
+            name: nil
+
+  @type t(state) :: %__MODULE__{
+          state: state,
+          name: String.t() | nil
+        }
+end
+```
+
 ### Documentation
 
 To add a `@typedoc` to the struct type, just add the attribute in the

lib/typed_struct.ex

@@ -8,6 +8,7 @@ defmodule TypedStruct do
   @accumulating_attrs [
     :ts_plugins,
     :ts_plugin_fields,
+    :ts_parameters,
     :ts_fields,
     :ts_types,
     :ts_enforce_keys
@@ -74,6 +75,18 @@ defmodule TypedStruct do
           field :field_four, atom(), default: :hey
         end
       end
+
+  You can also add type parameters:
+
+      defmodule MyModule do
+        use TypedStruct
+
+        typedstruct do
+          parameter :parameter
+
+          field :field, parameter
+        end
+      end
   """
   defmacro typedstruct(opts \\ [], do: block) do
     ast = TypedStruct.__typedstruct__(block, opts)
@@ -110,19 +123,23 @@ defmodule TypedStruct do
       @enforce_keys @ts_enforce_keys
       defstruct @ts_fields
 
-      TypedStruct.__type__(@ts_types, unquote(opts))
+      TypedStruct.__type__(@ts_parameters, @ts_types, unquote(opts))
     end
   end
 
   @doc false
-  defmacro __type__(types, opts) do
+  defmacro __type__(parameters, types, opts) do
     if Keyword.get(opts, :opaque, false) do
-      quote bind_quoted: [types: types] do
-        @opaque t() :: %__MODULE__{unquote_splicing(types)}
+      quote bind_quoted: [parameters: parameters, types: types] do
+        @opaque t(unquote_splicing(parameters)) :: %__MODULE__{
+                  unquote_splicing(types)
+                }
       end
     else
-      quote bind_quoted: [types: types] do
-        @type t() :: %__MODULE__{unquote_splicing(types)}
+      quote bind_quoted: [parameters: parameters, types: types] do
+        @type t(unquote_splicing(parameters)) :: %__MODULE__{
+                unquote_splicing(types)
+              }
       end
     end
   end
@@ -155,6 +172,22 @@ defmodule TypedStruct do
     end
   end
 
+  @doc """
+  Defines a type parameter for a typed struct.
+
+  ## Example
+
+      # A type parameter named int
+      parameter :int
+
+      field :number, int # not int()
+  """
+  defmacro parameter(name) do
+    quote bind_quoted: [name: name] do
+      TypedStruct.__parameter__(name, __ENV__)
+    end
+  end
+
   @doc """
   Defines a field in a typed struct.
 
@@ -175,6 +208,16 @@ defmodule TypedStruct do
     end
   end
 
+  @doc false
+  def __parameter__(name, %Macro.Env{module: mod}) when is_atom(name) do
+    Module.put_attribute(mod, :ts_parameters, Macro.var(name, mod))
+  end
+
+  def __parameter__(name, _env) do
+    raise ArgumentError,
+          "a parameter name must be an atom, got #{inspect(name)}"
+  end
+
   @doc false
   def __field__(name, type, opts, %Macro.Env{module: mod} = env)
       when is_atom(name) do

test/typed_struct_test.exs

@@ -20,6 +20,21 @@ defmodule TypedStructTest do
       def enforce_keys, do: @enforce_keys
     end
 
+  {:module, _name, bytecode_parameters, _exports} =
+    defmodule ParameterTestStruct do
+      use TypedStruct
+
+      @typep int() :: integer()
+
+      typedstruct do
+        parameter :str
+
+        field :int, int()
+        field :string, str
+        field :mandatory_string, str, enforce: true
+      end
+    end
+
   {:module, _name, bytecode_opaque, _exports} =
     defmodule OpaqueTestStruct do
       use TypedStruct
@@ -61,6 +76,7 @@ defmodule TypedStructTest do
     end
 
   @bytecode bytecode
+  @bytecode_parameters bytecode_parameters
   @bytecode_opaque bytecode_opaque
   @bytecode_noalias bytecode_noalias
 
@@ -130,6 +146,38 @@ defmodule TypedStructTest do
     assert type1 == type2
   end
 
+  test "generates a parameterized type for the struct", context do
+    # Define a second struct with the type expected for ParameterTestStruct.
+    {:module, _name, bytecode2, _exports} =
+      defmodule ParameterTestStruct2 do
+        @typep int() :: integer()
+
+        @enforce_keys [:mandatory_string]
+
+        defstruct [:int, :string, :mandatory_string]
+
+        @type t(str) :: %__MODULE__{
+                int: int() | nil,
+                string: str | nil,
+                mandatory_string: str
+              }
+      end
+
+    # Get both types and standardise them (remove line numbers and rename
+    # the second struct with the name of the first one).
+    type1 =
+      @bytecode_parameters
+      |> extract_first_type()
+      |> standardise(TypedStructTest.ParameterTestStruct)
+
+    type2 =
+      bytecode2
+      |> extract_first_type()
+      |> standardise(TypedStructTest.ParameterTestStruct2)
+
+    assert type1 == type2
+  end
+
   test "generates an opaque type if `opaque: true` is set" do
     # Define a second struct with the type expected for TestStruct.
     {:module, _name, bytecode_expected, _exports} =
@@ -248,20 +296,29 @@ defmodule TypedStructTest do
   defp standardise(type_info, struct \\ @standard_struct_name)
 
   defp standardise({name, type, params}, struct) when is_tuple(type),
-    do: {name, standardise(type, struct), params}
+    do: {name, standardise(type, struct), standardise_params(params)}
 
   defp standardise({:type, _, type, params}, struct),
     do: {:type, :line, type, standardise(params, struct)}
 
+  defp standardise({:user_type, _, type, params}, struct),
+    do: {:user_type, :line, type, standardise(params, struct)}
+
   defp standardise({:remote_type, _, params}, struct),
     do: {:remote_type, :line, standardise(params, struct)}
 
   defp standardise({:atom, _, struct}, struct),
     do: {:atom, :line, @standard_struct_name}
 
-  defp standardise({type, _, litteral}, _struct),
-    do: {type, :line, litteral}
+  defp standardise({type, _, literal}, _struct),
+    do: {type, :line, literal}
 
   defp standardise(list, struct) when is_list(list),
     do: Enum.map(list, &standardise(&1, struct))
+
+  defp standardise_params(params) when is_list(params),
+    do: Enum.map(params, &standardise_params/1)
+
+  defp standardise_params({:var, _, name}), do: {:var, :id, name}
+  defp standardise_params(params), do: params
 end