fixed #221 (#222)

Author: hrshdhgd

sssom/cli.py

@@ -232,16 +232,22 @@ def dosql(query: str, inputs: List[str], output: TextIO):
     """  # noqa: DAR101
     # should start with from_tsv and MOST should return write_sssom
     n = 1
+    new_msdf = MappingSetDataFrame()
     while len(inputs) >= n:
         fn = inputs[n - 1]
-        df = read_sssom_table(fn).df
+        msdf = read_sssom_table(fn)
+        df = msdf.df
         # df = parse(fn)
         globals()[f"df{n}"] = df
         tn = re.sub("[.].*", "", Path(fn).stem).lower()
         globals()[tn] = df
         n += 1
-    df = sqldf(query)
-    df.to_csv(output, sep="\t", index=False)
+
+    new_df = sqldf(query)
+    new_msdf.df = new_df
+    new_msdf.prefix_map = msdf.prefix_map
+    new_msdf.metadata = msdf.metadata
+    write_table(new_msdf, output)
 
 
 @main.command()

tests/test_cli.py

@@ -13,6 +13,7 @@
     crosstab,
     dedupe,
     diff,
+    dosql,
     merge,
     parse,
     partition,
@@ -56,6 +57,8 @@ def test_cli_single_input(self):
                     self.run_crosstab(runner, test)
                     self.run_correlations(runner, test)
                     self.run_reconcile_prefix(runner, test)
+                    self.run_dosql(runner, test)
+
         self.assertTrue(len(test_cases) > 2)
 
     def test_cli_multiple_input(self):
@@ -70,7 +73,6 @@ def test_cli_multiple_input(self):
 
     def run_successful(self, result: Result, test_case: SSSOMTestCase) -> None:
         """Check the test result is successful."""
-        # self.assertTrue(result.exit_code == 0, f"Run failed with message {result.exception}")
         self.assertEqual(
             result.exit_code,
             0,
@@ -249,3 +251,19 @@ def run_reconcile_prefix(
         )
         self.run_successful(result, test_case)
         return result
+
+    def run_dosql(self, runner: CliRunner, test_case: SSSOMTestCase) -> Result:
+        """Test a simple dosql command."""
+        out_file = os.path.join(test_out_dir, "dosql_test.tsv")
+        result = runner.invoke(
+            dosql,
+            [
+                "-q",
+                "SELECT * FROM df WHERE subject_label = 'heart'",
+                test_case.filepath,
+                "-o",
+                os.path.join(test_out_dir, out_file),
+            ],
+        )
+        self.run_successful(result, test_case)
+        return result