From 7409fd11a6fb09f4dbd5f51497c835dcd5f45c02 Mon Sep 17 00:00:00 2001
From: Rok Mandeljc <rok.mandeljc@gmail.com>
Date: Sat, 16 Dec 2023 19:02:27 +0100
Subject: [PATCH] hooks: add hook for fvcore.nn

Add hook for `fvcore.nn` to collect its source .py files for
torchscript/JIT.

Add a basic import test that demonstrates the need for that.
---
 news/676.new.10.rst                                |  2 ++
 .../hooks/stdhooks/hook-fvcore.nn.py               | 14 ++++++++++++++
 .../tests/test_deep_learning.py                    |  9 +++++++++
 3 files changed, 25 insertions(+)
 create mode 100644 news/676.new.10.rst
 create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-fvcore.nn.py

diff --git a/news/676.new.10.rst b/news/676.new.10.rst
new file mode 100644
index 000000000..376795b34
--- /dev/null
+++ b/news/676.new.10.rst
@@ -0,0 +1,2 @@
+Add hook for ``fvcore.nn`` to collect its source .py files for
+TorchScript/JIT.
diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-fvcore.nn.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-fvcore.nn.py
new file mode 100644
index 000000000..6f47b491c
--- /dev/null
+++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-fvcore.nn.py
@@ -0,0 +1,14 @@
+# ------------------------------------------------------------------
+# Copyright (c) 2023 PyInstaller Development Team.
+#
+# This file is distributed under the terms of the GNU General Public
+# License (version 2.0 or later).
+#
+# The full license is available in LICENSE.GPL.txt, distributed with
+# this software.
+#
+# SPDX-License-Identifier: GPL-2.0-or-later
+# ------------------------------------------------------------------
+
+# Collect source .py files for JIT/torchscript. Requires PyInstaller >= 5.3, no-op in older versions.
+module_collection_mode = 'pyz+py'
diff --git a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py
index f462d7eb4..4e54d556f 100644
--- a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py
+++ b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py
@@ -278,3 +278,12 @@ def forward(self, x):
         print("Test X:", test_x.numpy())
         print("Predicted Y:", observed_pred.mean.numpy())
     """)
+
+
+# Basic import test for fvcore.nn, which shows that we need to collect its source.py files for TorchScript/JIT.
+@importorskip('fvcore')
+@onedir_only
+def test_fvcore(pyi_builder):
+    pyi_builder.test_source("""
+        import fvcore.nn
+    """)