From 51a30ede95f15297a5c4fe780e5de1efd225ba42 Mon Sep 17 00:00:00 2001
From: Sebastien Mondet <sebastien.mondet@ahrefs.com>
Date: Wed, 27 Nov 2024 15:47:10 -0500
Subject: [PATCH] Add option `--only-matching` to limit output

---
 bin/main.ml      | 17 +++++++++++++++--
 lib/generator.ml | 10 +++++++++-
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/bin/main.ml b/bin/main.ml
index 860a534..2ab2160 100644
--- a/bin/main.ml
+++ b/bin/main.ml
@@ -46,11 +46,24 @@ let main =
   let doc = "Generate an ATD file from a list of JSON Schema / OpenAPI document" in
   let state_term =
     Term.(
-      const (fun skip_doc pad ->
-        Generator.{ with_doc = not skip_doc; protect_against_duplicates = (if pad then Some (ref []) else None) }
+      const (fun skip_doc pad toplevel_types ->
+        Generator.
+          {
+            with_doc = not skip_doc;
+            protect_against_duplicates = (if pad then Some (ref []) else None);
+            toplevel_types;
+          }
       )
       $ Arg.(value (flag (info [ "skip-doc" ] ~doc:"Skip documentation annotations.")))
       $ Arg.(value (flag (info [ "protect-against-duplicates" ] ~doc:"Make sure no duplicate types are generated.")))
+      $ (const (function
+           | [] -> `All
+           | some -> `Only some
+           )
+        $ Arg.(
+            value (opt_all string [] (info [ "only-matching" ] ~docv:"REGEXP" ~doc:"Only output types matching REGEXP."))
+          )
+        )
     )
   in
   let paths = Arg.(non_empty & pos_all file [] & info [] ~docv:"FILES" ~doc) in
diff --git a/lib/generator.ml b/lib/generator.ml
index 9542c69..3efa414 100644
--- a/lib/generator.ml
+++ b/lib/generator.ml
@@ -5,9 +5,10 @@ open Utils
 type state = {
   with_doc : bool;
   protect_against_duplicates : string list ref option;
+  toplevel_types : [ `All | `Only of string list ];
 }
 
-let default_state = { with_doc = true; protect_against_duplicates = None }
+let default_state = { with_doc = true; protect_against_duplicates = None; toplevel_types = `All }
 
 let record_field_name _state str =
   let cleaned_field_name = Utils.sanitize_name str in
@@ -288,6 +289,13 @@ type int64 = int <ocaml repr="int64">
     from
 
 let make_atd_of_schemas state schemas =
+  let schemas =
+    match state.toplevel_types with
+    | `All -> schemas
+    | `Only l ->
+      let res = List.map (Printf.ksprintf Str.regexp "^%s$") l in
+      List.filter (fun (name, _) -> List.exists (fun re -> Str.string_match re name 0) res) schemas
+  in
   input_toplevel_schemas :=
     List.filter_map
       (function