From a9ae6239a36cab5dfc5cc02f79b3792122dfe55e Mon Sep 17 00:00:00 2001
From: leej3 <johnleenimh@gmail.com>
Date: Fri, 14 May 2021 15:53:27 +0100
Subject: [PATCH] add additional support for transformation parsing

Add tests and add support for:
Model jsons with transformations
Transformation jsons
Python in memory representations of the above
---
 bids/__init__.py                |  1 -
 bids/variables/io.py            | 10 +++++-----
 bids/variables/tests/test_io.py | 35 +++++++++++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/bids/__init__.py b/bids/__init__.py
index 3d38ffc26..fb96d2d4b 100644
--- a/bids/__init__.py
+++ b/bids/__init__.py
@@ -14,7 +14,6 @@
     "reports",
     "utils",
     "variables",
-    "statsmodels_design_synthesizer",
 ]
 
 due.cite(Doi("10.1038/sdata.2016.44"),
diff --git a/bids/variables/io.py b/bids/variables/io.py
index 0a382d8d0..76842583d 100644
--- a/bids/variables/io.py
+++ b/bids/variables/io.py
@@ -584,13 +584,13 @@ def parse_transforms(transforms_in, validate=True,level="run"):
     # transformations has been obtained. This will most likely be the case since
     # transformations at higher levels will no longer be required when the new
     # "flow" approach is used.
-    if "nodes" in transforms_raw:
-        nodes_key = "nodes"
-    elif "steps" in transforms_raw:
-        nodes_key = "steps"
+    if "transformations" in transforms_raw:
+        transforms = transforms_raw["transformations"]
+    elif any(k in transforms_raw for k in ["nodes","steps"]):
+        nodes_key = "nodes" if "nodes" in transforms_raw else "steps"
+        transforms = transforms_raw[nodes_key][0]["transformations"]
     else:
         raise ValueError("Cannot find a key for nodes in the json input representing the model")
-    transforms = transforms_raw[nodes_key][0]["transformations"]
     return transforms
 
 
diff --git a/bids/variables/tests/test_io.py b/bids/variables/tests/test_io.py
index c4690ce84..ce2df52a1 100644
--- a/bids/variables/tests/test_io.py
+++ b/bids/variables/tests/test_io.py
@@ -2,12 +2,21 @@
 from bids.variables import (SparseRunVariable, SimpleVariable,
                             DenseRunVariable, load_variables)
 from bids.variables.entities import Node, RunNode, NodeIndex
+from bids.variables.io import parse_transforms
 from unittest.mock import patch
 import pytest
 from os.path import join
+from pathlib import Path
+import tempfile
+import json
 from bids.tests import get_test_data_path
 from bids.config import set_option, get_option
 
+EXAMPLE_TRANSFORM = {
+    "Transformations":[{"Name":"example_trans","Inputs":["col_a","col_b"]}]
+}
+TRANSFORMS_JSON = join(tempfile.tempdir,"tranformations.json")
+Path(TRANSFORMS_JSON).write_text(json.dumps(EXAMPLE_TRANSFORM))
 
 @pytest.fixture
 def layout1():
@@ -103,3 +112,29 @@ def test_load_synthetic_dataset(synthetic):
     subs = index.get_nodes('subject')
     assert len(subs) == 5
     assert set(subs[0].variables.keys()) == {'systolic_blood_pressure'}
+
+@pytest.mark.parametrize(
+    "test_case,transform_input,expected_names",
+    [
+        ("raw transform json",
+         EXAMPLE_TRANSFORM,
+         ["example_trans"]
+        ),
+        ("transform json file",
+         TRANSFORMS_JSON,
+         ["example_trans"]
+        ),
+        ("raw model json",
+         {"Nodes": [EXAMPLE_TRANSFORM]},
+         ["example_trans"]
+        ),
+         ("model json file",
+         str(Path(get_test_data_path()) / "ds005/models/ds-005_type-mfx_model.json"),
+         ["Scale"]
+        ),
+    ]
+)
+def test_parse_transforms(test_case,transform_input,expected_names):
+    result = parse_transforms(transform_input)
+    transformation_names =  [x['name'] for x in result]
+    assert expected_names == transformation_names