Commit 0f642753 authored by Asitha Senanayake's avatar Asitha Senanayake
Browse files

feat(#143): add 'only_groups' option to AGS4_to_dataframe() 

This feature allows a subset of the groups to be converted instead of
converting all the groups in the input file.
parent 40492868

python_ags4/AGS4.py

+12 −3
Original line number Diff line number Diff line
@@ -203,7 +203,8 @@ def AGS4_to_dict(filepath_or_buffer, encoding='utf-8', get_line_numbers=False, r 
    return data, headings





def AGS4_to_dataframe(filepath_or_buffer, encoding='utf-8', get_line_numbers=False, rename_duplicate_headers=True):

def AGS4_to_dataframe(filepath_or_buffer, encoding='utf-8', get_line_numbers=False, rename_duplicate_headers=True,

                      only_groups=None):

    """Load all the tables in an AGS4 file to a dictionary of Pandas dataframes.



    The output is a dictionary of dataframes with the name of each AGS4 table
@@ -223,6 +224,9 @@ def AGS4_to_dataframe(filepath_or_buffer, encoding='utf-8', get_line_numbers=Fal 
        Rename duplicate headers if found. Neither AGS4 tables nor Pandas

        dataframes allow duplicate headers, therefore a number will be appended

        to duplicates to make them unique.

    only_groups : list or None (default=None)

        An optional list of groups to convert instead of converting all the

        groups in the input file.



    Returns

    -------
@@ -263,6 +267,11 @@ def AGS4_to_dataframe(filepath_or_buffer, encoding='utf-8', get_line_numbers=Fal 


    # Convert dictionary of dictionaries to a dictionary of Pandas dataframes

    tables = {}



    if only_groups:

        for key in only_groups:

            tables[key] = DataFrame(data[key])

    else:

        for key in data:

            tables[key] = DataFrame(data[key])

tests/test_ags4.py

+6 −2
Original line number Diff line number Diff line
@@ -64,12 +64,16 @@ def test_AGS4_bytestream_to_dict(LOCA=LOCA): 
    assert tables['LOCA'] == LOCA





def test_AGS4_file_to_dataframe(LOCA=LOCA):

    tables, headings = AGS4.AGS4_to_dataframe(TEST_DATA)

@pytest.mark.parametrize("only_groups", [None, ['PROJ', 'TRAN', 'LOCA']])

def test_AGS4_file_to_dataframe(only_groups, LOCA=LOCA):

    tables, headings = AGS4.AGS4_to_dataframe(TEST_DATA, only_groups=only_groups)



    assert tables['LOCA'].loc[2, 'LOCA_ID'] == 'Location_1'

    assert tables['LOCA'].equals(pd.DataFrame(LOCA))



    if only_groups:

        assert list(tables.keys()) == only_groups





def test_AGS4_stream_to_dataframe(LOCA=LOCA):

    with open(TEST_DATA, 'r') as file: