diff --git a/q2_amrfinderplus/types/_transformer.py b/q2_amrfinderplus/types/_transformer.py index 06b6900..3c2ad29 100644 --- a/q2_amrfinderplus/types/_transformer.py +++ b/q2_amrfinderplus/types/_transformer.py @@ -37,11 +37,11 @@ def _transfomer_helper(data): else: id_value = file_dir_name + "/" + file_name[:-18] - create_append_df( - file_path=file, - df_list=df_list, - id_value=id_value, - ) + # Create df and append it to df_list + df = pd.read_csv(file, sep="\t") + df.insert(0, "Sample/MAG_ID", id_value) + df_list.append(df) + else: # Annotations file from feature data mags or sample data contigs if file_dir_name.endswith("_amr_annotations.tsv"): @@ -50,26 +50,14 @@ def _transfomer_helper(data): else: id_value = file_dir_name[:-22] - create_append_df( - file_path=os.path.join(str(data), file_dir_name), - df_list=df_list, - id_value=id_value, - ) + # Create df and append it to df_list + df = pd.read_csv(os.path.join(str(data), file_dir_name), sep="\t") + df.insert(0, "Sample/MAG_ID", id_value) + df_list.append(df) return combine_dataframes(df_list) -def create_append_df(file_path, df_list, id_value): - # Read in df - df = pd.read_csv(file_path, sep="\t") - - # Insert column with sample or mag IDs - df.insert(0, "Sample/MAG_ID", id_value) - - # Append df to df list - df_list.append(df) - - def combine_dataframes(df_list): # Concat all dfs df_combined = pd.concat(df_list, axis=0) diff --git a/q2_amrfinderplus/types/tests/test_types_formats_transformers.py b/q2_amrfinderplus/types/tests/test_types_formats_transformers.py index 5120733..34bcfb1 100644 --- a/q2_amrfinderplus/types/tests/test_types_formats_transformers.py +++ b/q2_amrfinderplus/types/tests/test_types_formats_transformers.py @@ -7,7 +7,6 @@ # ---------------------------------------------------------------------------- import os import tempfile -from io import StringIO import pandas as pd import qiime2 @@ -20,11 +19,7 @@ AMRFinderPlusAnnotationsDirFmt, AMRFinderPlusDatabaseDirFmt, ) -from q2_amrfinderplus.types._transformer import ( - _transfomer_helper, - combine_dataframes, - create_append_df, -) +from q2_amrfinderplus.types._transformer import _transfomer_helper, combine_dataframes class TestAMRFinderPlusTypesAndFormats(TestPluginBase): @@ -125,20 +120,12 @@ def test_amrfinderplus_annotations_dir_fmt_path_maker(self): self.assertEqual(str(path), os.path.join(str(fmt), "id_amr_annotations.tsv")) -class MetadataUtilsTest(TestPluginBase): +class MetadataTransformerUtilsTest(TestPluginBase): package = "q2_amrfinderplus.types.tests" def setUp(self): super().setUp() # Setup test data - self.file_data_1 = "col1\tcol2\nval1\tval2\nval3\tval4" - self.file_data_2 = "col1\tcol2\nval5\tval6\nval7\tval8" - - self.df1 = pd.read_csv(StringIO(self.file_data_1), sep="\t") - self.df2 = pd.read_csv(StringIO(self.file_data_2), sep="\t") - - self.df_list = [] - self.df1 = pd.DataFrame( { "Sample/MAG_ID": ["id_value_1", "id_value_1"], @@ -171,14 +158,6 @@ def setUp(self): self.expected_combined_df.index = self.expected_combined_df.index.astype(str) self.expected_combined_df.index.name = "id" - def test_create_append_df(self): - # Test create_append_df function - create_append_df(StringIO(self.file_data_1), self.df_list, "id_value_1") - create_append_df(StringIO(self.file_data_2), self.df_list, "id_value_2") - - pd.testing.assert_frame_equal(self.df_list[0], self.df1) - pd.testing.assert_frame_equal(self.df_list[1], self.df2) - def test_combine_dataframes(self): # Test combine_dataframes function df_list = [self.df1, self.df2]