You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
59 lines
2.3 KiB
59 lines
2.3 KiB
from pandas import DataFrame
|
|
import re
|
|
|
|
|
|
COLUMN_NAME_REGEX = re.compile(r"(?P<column_name>(\w|\.|#|\/)+)", re.IGNORECASE)
|
|
|
|
def replace_bad_cols(line: str, cols: list[str]):
|
|
"""
|
|
Replaces bad column names in a string with modified names that have spaces replaced with dots.
|
|
|
|
Args:
|
|
line (str): The string containing the column names to modify.
|
|
cols (list[str]): A list of column names to modify.
|
|
|
|
Returns:
|
|
str: The modified string with bad column names replaced.
|
|
"""
|
|
for c in cols:
|
|
# Replace spaces with dots in the column name
|
|
gc = c.replace(' ', '.')
|
|
# Replace the bad column name with the modified column name in the string
|
|
line = line.replace(c, gc)
|
|
return line
|
|
|
|
|
|
def extract_data(input_doc: str, column_list: list[str]):
|
|
"""
|
|
Extracts data from a string in a table-like format, where columns are identified by a list of column names, and
|
|
returns the data as a Pandas DataFrame.
|
|
|
|
Args:
|
|
input_doc (str): The string containing the table-like data to extract.
|
|
column_list (list[str]): A list of column names to identify the columns in the table-like data.
|
|
|
|
Returns:
|
|
pandas.DataFrame: A DataFrame containing the extracted data from the input string.
|
|
"""
|
|
line: str
|
|
columns = {}
|
|
data = {}
|
|
for line in input_doc.splitlines():
|
|
if len(columns) == 0 :
|
|
# Find the line that contains the column names and replace bad column names
|
|
if re.search("^\w", line):
|
|
line = replace_bad_cols(line, column_list)
|
|
# Find the start and end positions of each column name and store them in a dictionary
|
|
columns_names = re.finditer(COLUMN_NAME_REGEX, line)
|
|
for c in columns_names:
|
|
columns[c.group("column_name")] = {"start": c.start(), "end": c.end()}
|
|
data[c.group("column_name")] = []
|
|
continue
|
|
elif len(line) < 2:
|
|
continue
|
|
# Check if we've reached the end of the table and return the data
|
|
if re.search("\d+ records listed", line):
|
|
return DataFrame(data)
|
|
# Extract the data from each column based on the start and end positions
|
|
for key, span in columns.items():
|
|
data[key].append(line[span["start"]:span["end"]].strip()) |