co-analyst/analyst.py at main · Mittal-Analytics/co-analyst · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
import json
import os

from utilities import explorer, grader, tablex, tools


def _extract_data_from_table(statement_name, table, grading, column_names, unit):
    stack = [{"title": statement_name, "unit": unit, "data": []}]
    for row in table:
        try:
            while (
                len(stack) > 1
                and grading[row[0]["title"]] >= grading[stack[-1]["title"]]
            ):
                stack[-1]["title"] = tools.remove_list_marker(stack[-1]["title"])
                stack[-2]["data"].append(stack[-1])
                stack.pop()
        except KeyError:
            if len(row) == 2:
                stack[-1]["title"] = tools.remove_list_marker(stack[-1]["title"])
                stack[-2]["data"].append(stack[-1])
                stack.pop()
        stack.append({"title": row[0]["title"], "data": []})
        if len(row) == 2:
            stack[-1]["title"] = stack[-2]["title"]
            stack[-1]["data"].append(
                {
                    column_names[1]: tools.negate(
                        tools.remove_list_marker(row[0]["title"])
                    ),
                    column_names[2]: tools.negate(
                        tools.remove_list_marker(row[1]["title"])
                    ),
                }
            )
        elif len(row) == 3:
            stack[-1]["data"].append(
                {
                    column_names[1]: tools.negate(
                        tools.remove_list_marker(row[1]["title"])
                    ),
                    column_names[2]: tools.negate(
                        tools.remove_list_marker(row[2]["title"])
                    ),
                }
            )
        elif len(row) == 4:
            stack[-1]["data"].append(
                {
                    column_names[0]: tools.remove_list_marker(row[1]["title"]),
                    column_names[1]: tools.negate(
                        tools.remove_list_marker(row[2]["title"])
                    ),
                    column_names[2]: tools.negate(
                        tools.remove_list_marker(row[3]["title"])
                    ),
                }
            )

    while len(stack) > 1:
        stack[-1]["title"] = tools.remove_list_marker(stack[-1]["title"])
        stack[-2]["data"].append(stack[-1])
        stack.pop()

    return stack[0]


def extract_data_from_pdf(pdf_path, **kwargs):
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"File not found at {pdf_path}.")

    try:
        start, end = kwargs["start"], kwargs["end"]
        statement_name = explorer.find_statement_name(pdf_path, start)
    except KeyError:
        try:
            statement_name = kwargs["statement_name"].lower()
            start, end = explorer.find_page_range(pdf_path, statement_name)
        except KeyError:
            raise KeyError("Either start and end or statement_name must be provided.")

    unit = explorer.find_unit(pdf_path, start)
    tables = []
    for page_num in range(start - 1, end):
        tables.append(tablex.extract_tables(pdf_path, page_num))

    response = []
    for table in tables:
        first_row = explorer.find_first_row(table)
        column_names = explorer.find_column_names(table[0:first_row])
        table = table[first_row:]

        grading = grader.get_grading(table)
        data = _extract_data_from_table(
            statement_name, table, grading, column_names, unit
        )

        # data["title"] will be None when table on the current page
        # is continuation of the table of the previous page.
        if data["title"] is None:
            response[-1]["data"].append(data["data"][0])
        else:
            response.append(data)
    return json.dumps(response)