From cfa37b2078e0503eaccd7ce20cc287dbc4859b47 Mon Sep 17 00:00:00 2001
From: Filippo Bigi <98903385+frostedoyster@users.noreply.github.com>
Date: Fri, 26 Jan 2024 12:54:57 +0100
Subject: [PATCH] Display `metatensor-torch` torchscript classes correctly in
 docs (#40)

* Display torchscript classes correctly in docs
* Disable compiled functions when building docs
---
 docs/src/conf.py                            |  4 ++++
 src/metatensor/models/utils/data/dataset.py | 10 ++++++++--
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/docs/src/conf.py b/docs/src/conf.py
index 453bb1047..ecdee45d7 100644
--- a/docs/src/conf.py
+++ b/docs/src/conf.py
@@ -10,6 +10,10 @@
 ROOT = os.path.abspath(os.path.join("..", ".."))
 sys.path.insert(0, ROOT)
 
+# when importing metatensor-torch, this will change the definition of the classes
+# to include the documentation
+os.environ["METATENSOR_IMPORT_FOR_SPHINX"] = "1"
+
 # -- Project information -----------------------------------------------------
 
 # The master toctree document.
diff --git a/src/metatensor/models/utils/data/dataset.py b/src/metatensor/models/utils/data/dataset.py
index b0144a5bd..6de8e1c1a 100644
--- a/src/metatensor/models/utils/data/dataset.py
+++ b/src/metatensor/models/utils/data/dataset.py
@@ -1,3 +1,4 @@
+import os
 from typing import Dict, List
 
 import metatensor.torch
@@ -6,8 +7,13 @@
 from metatensor.torch.atomistic import ModelCapabilities, System
 
 
-compiled_slice = torch.jit.script(metatensor.torch.slice)
-compiled_join = torch.jit.script(metatensor.torch.join)
+if os.environ.get("METATENSOR_IMPORT_FOR_SPHINX", "0") == "1":
+    # This is necessary to make the Sphinx documentation build
+    compiled_slice = None
+    compiled_join = None
+else:
+    compiled_slice = torch.jit.script(metatensor.torch.slice)
+    compiled_join = torch.jit.script(metatensor.torch.join)
 
 
 class Dataset(torch.utils.data.Dataset):