diff --git a/docs/example_ectd2.ipynb b/docs/example_ectd2.ipynb index 8c00df5f..4916ab3f 100644 --- a/docs/example_ectd2.ipynb +++ b/docs/example_ectd2.ipynb @@ -172,11 +172,6 @@ "To include the upper CTOC sections that include search results:" ] }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [] - }, { "cell_type": "code", "execution_count": null, @@ -207,7 +202,7 @@ }, "outputs": [], "source": [ - "print(\"\\n\".join(ctoc.find_section_given_words([\"REMS\",\"isE\"], outfmt='tree')))" + "print(\"\\n\".join(ctoc.find_section_given_words(words=[\"REMS\",\"isE\"], outfmt='tree')))" ] }, { diff --git a/mtbp3/health/ectd.py b/mtbp3/health/ectd.py index e5ed1022..2e91b199 100644 --- a/mtbp3/health/ectd.py +++ b/mtbp3/health/ectd.py @@ -48,12 +48,17 @@ def show_ctoc_tree(self, module=None, to_right=False): return tree.list_tree(to_right=to_right) def find_section_given_words(self, words, outfmt='simple', include='up', to_right=False): - if not isinstance(words, str) or not words: - raise ValueError("str must be a nonempty string") + if isinstance(words, str) and words: + words = [words] + elif isinstance(words, list) and len(words)>0: + assert all(isinstance(word, str) for word in words), "Elements in the list must be strings" + else: + raise ValueError("words must be a string or a list") + if include not in ['up', 'both']: raise ValueError("Invalid value for include. Supported values are 'up' and 'both'.") - out = [row for row in self.ctoc if words.lower() in row.lower()] + out = [row for row in self.ctoc if any(word.lower() in row.lower() for word in words)] if outfmt == 'simple': return out elif outfmt == 'tree': diff --git a/tests/test_health_ectd.py b/tests/test_health_ectd.py new file mode 100644 index 00000000..ff27cfb7 --- /dev/null +++ b/tests/test_health_ectd.py @@ -0,0 +1,32 @@ +import unittest +from mtbp3.health.ectd import ctoc_by_fda + +class TestCtocByFDA(unittest.TestCase): + + def setUp(self): + self.ctoc = ctoc_by_fda(ectd_version="3.2.2", ctoc_version="2.3.3") + + def test_load_list(self): + self.assertIsInstance(self.ctoc.ctoc, list) + self.assertGreater(len(self.ctoc.ctoc), 0) + + def test_show_ctoc_tree(self): + tree = self.ctoc.show_ctoc_tree(module=1, to_right=False) + self.assertIsInstance(tree[0], str) + self.assertGreater(len(tree), 0) + + def test_find_section_given_words(self): + result = self.ctoc.find_section_given_words("REMS", outfmt='simple', include='up', to_right=False) + self.assertIsInstance(result, list) + self.assertGreater(len(result), 0) + + def test_find_section_given_words_invalid_include(self): + with self.assertRaises(ValueError): + self.ctoc.find_section_given_words("REMS", outfmt='simple', include='invalid', to_right=False) + + def test_find_section_given_words_invalid_outfmt(self): + with self.assertRaises(ValueError): + self.ctoc.find_section_given_words("REMS", outfmt='invalid', include='up', to_right=False) + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/tests/test_emt.py b/tests/test_health_emt.py similarity index 100% rename from tests/test_emt.py rename to tests/test_health_emt.py