Add python version of findExampleDataFile

python/CMakeLists.txt

@@ -4,16 +4,9 @@ if (NOT GTSAM_BUILD_PYTHON)
     return()
 endif()
 
-# Common directory for storing data stored with the package.
+# Common directory for storing data/datasets stored with the package.
-# This is because the exact location of the Python installation may vary
+# This will store the data in the Python site package directly.
-# depending on various factors such as python version and use of virtual envs.
+set(GTSAM_PYTHON_DATASET_DIR "./gtsam/Data")
-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()
 
 # Generate setup.py.
 file(READ "${PROJECT_SOURCE_DIR}/README.md" README_CONTENTS)

python/gtsam/__init__.py

@@ -1,4 +1,6 @@
+from . import utils
 from .gtsam import *
+from .utils import findExampleDataFile
 
 
 def _init():

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