From 43dfe0f47d8e001572b7cb2471f5ca90515e2dfd Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Mon, 31 Aug 2020 13:56:56 -0400
Subject: [PATCH] Add python version of findExampleDataFile

---
 python/CMakeLists.txt          | 13 +++----------
 python/gtsam/__init__.py       |  2 ++
 python/gtsam/utils/__init__.py | 22 ++++++++++++++++++++++
 3 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
index 6fad7ba45..e15e571f5 100644
--- a/python/CMakeLists.txt
+++ b/python/CMakeLists.txt
@@ -4,16 +4,9 @@ if (NOT GTSAM_BUILD_PYTHON)
     return()
 endif()
 
-# Common directory for storing data stored with the package.
-# This is because the exact location of the Python installation may vary
-# depending on various factors such as python version and use of virtual envs.
-if(CMAKE_SYSTEM_NAME STREQUAL "Linux" OR CMAKE_SYSTEM_NAME STREQUAL "Darwin")
-    # Easy directory with write access.
-    set(GTSAM_PYTHON_DATASET_DIR "/tmp/gtsam_example_data")
-elseif(CMAKE_SYSTEM_NAME STREQUAL "Windows")
-    # %LocalAppData% directory in Windows.
-    set(GTSAM_PYTHON_DATASET_DIR "$ENV{LOCALAPPDATA}/gtsam_example_data")
-endif()
+# Common directory for storing data/datasets stored with the package.
+# This will store the data in the Python site package directly.
+set(GTSAM_PYTHON_DATASET_DIR "./gtsam/Data")
 
 # Generate setup.py.
 file(READ "${PROJECT_SOURCE_DIR}/README.md" README_CONTENTS)
diff --git a/python/gtsam/__init__.py b/python/gtsam/__init__.py
index b2d1fb7e7..70a00c3dc 100644
--- a/python/gtsam/__init__.py
+++ b/python/gtsam/__init__.py
@@ -1,4 +1,6 @@
+from . import utils
 from .gtsam import *
+from .utils import findExampleDataFile
 
 
 def _init():
diff --git a/python/gtsam/utils/__init__.py b/python/gtsam/utils/__init__.py
index e69de29bb..5ee733ba4 100644
--- a/python/gtsam/utils/__init__.py
+++ b/python/gtsam/utils/__init__.py
@@ -0,0 +1,22 @@
+import glob
+import os.path as osp
+
+
+def findExampleDataFile(name):
+    """
+    Find the example data file specified by `name`.
+    """
+    # This is okay since gtsam will already be loaded
+    # before this function is accessed. Thus no effect.
+    import gtsam
+
+    site_package_path = gtsam.__path__[0]
+    # add the * at the end to glob the entire directory
+    data_path = osp.join(site_package_path, "Data", "*")
+
+    extensions = ("", ".graph", ".txt", ".out", ".xml", ".g2o")
+
+    for data_file_path in glob.glob(data_path, recursive=True):
+        for ext in extensions:
+            if (name + ext) == osp.basename(data_file_path):
+                return data_file_path