Merge pull request #501 from borglab/feature/python-data

Example Datasets with Python Package

gtsam/gtsam.i

@@ -2800,7 +2800,6 @@ class SfmData {
   gtsam::SfmTrack track(size_t idx) const;
 };
 
-string findExampleDataFile(string name);
 pair<gtsam::NonlinearFactorGraph*, gtsam::Values*> load2D(string filename,
     gtsam::noiseModel::Diagonal* model, int maxIndex, bool addNoise, bool smart);
 pair<gtsam::NonlinearFactorGraph*, gtsam::Values*> load2D(string filename,

python/CMakeLists.txt

@@ -4,6 +4,10 @@ if (NOT GTSAM_BUILD_PYTHON)
     return()
 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)
 configure_file(${PROJECT_SOURCE_DIR}/python/setup.py.in

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

python/setup.py.in

@@ -1,3 +1,4 @@
+import glob
 import os
 import sys
 
@@ -8,6 +9,11 @@ except ImportError:
 
 packages = find_packages(where=".")
 print("PACKAGES: ", packages)
+
+data_path = '${GTSAM_SOURCE_DIR}/examples/Data/'
+data_files_and_directories = glob.glob(data_path + '**', recursive=True)
+data_files = [x for x in data_files_and_directories if not os.path.isdir(x)]
+
 package_data = {
     '': [
         './*.so',
@@ -44,6 +50,7 @@ setup(
     ],
     packages=packages,
     package_data=package_data,
+    data_files=[('${GTSAM_PYTHON_DATASET_DIR}', data_files),],
     test_suite="gtsam.tests",
     install_requires=["numpy"],
     zip_safe=False,