openbb-us-eia 1.0.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- openbb_us_eia-1.0.0/PKG-INFO +140 -0
- openbb_us_eia-1.0.0/README.md +119 -0
- openbb_us_eia-1.0.0/openbb_us_eia/__init__.py +25 -0
- openbb_us_eia-1.0.0/openbb_us_eia/models/__init__.py +1 -0
- openbb_us_eia-1.0.0/openbb_us_eia/models/petroleum_status_report.py +241 -0
- openbb_us_eia-1.0.0/openbb_us_eia/models/short_term_energy_outlook.py +284 -0
- openbb_us_eia-1.0.0/openbb_us_eia/utils/__init__.py +1 -0
- openbb_us_eia-1.0.0/openbb_us_eia/utils/constants.py +1272 -0
- openbb_us_eia-1.0.0/openbb_us_eia/utils/helpers.py +43 -0
- openbb_us_eia-1.0.0/openbb_us_eia/utils/py.typed +0 -0
- openbb_us_eia-1.0.0/openbb_us_eia/utils/total_energy_map.json +982 -0
- openbb_us_eia-1.0.0/pyproject.toml +22 -0
|
@@ -0,0 +1,140 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: openbb-us-eia
|
|
3
|
+
Version: 1.0.0
|
|
4
|
+
Summary: The U.S. Energy Information Administration is committed to its free and open data by making it available through an Application Programming Interface (API) and its open data tools. See https://www.eia.gov/opendata/ for more information.
|
|
5
|
+
License: AGPL-3.0-only
|
|
6
|
+
Author: OpenBB Team
|
|
7
|
+
Author-email: hello@openbb.co
|
|
8
|
+
Requires-Python: >=3.9,<4.0
|
|
9
|
+
Classifier: License :: OSI Approved :: GNU Affero General Public License v3
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
16
|
+
Requires-Dist: async-lru (>=2.0.4,<3.0.0)
|
|
17
|
+
Requires-Dist: openbb-core (>=1.3.6,<2.0.0)
|
|
18
|
+
Requires-Dist: openpyxl (>=3.1.5,<4.0.0)
|
|
19
|
+
Requires-Dist: xlrd (>=2.0.1,<3.0.0)
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
|
|
22
|
+
# OpenBB EIA Provider Extension
|
|
23
|
+
|
|
24
|
+
This module integrates the [EIA](https://eia.gov) data provider into the OpenBB Platform.
|
|
25
|
+
|
|
26
|
+
## Installation
|
|
27
|
+
|
|
28
|
+
### PyPI
|
|
29
|
+
|
|
30
|
+
```sh
|
|
31
|
+
pip install openbb-us-eia
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
### From Source
|
|
35
|
+
|
|
36
|
+
After cloning the main repository, navigate into this folder and enter:
|
|
37
|
+
|
|
38
|
+
```sh
|
|
39
|
+
pip install .
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
To install in editable mode:
|
|
43
|
+
|
|
44
|
+
```sh
|
|
45
|
+
pip install -e .
|
|
46
|
+
```
|
|
47
|
+
|
|
48
|
+
## Authorization
|
|
49
|
+
|
|
50
|
+
Functions calling the EIA's API require free registration and an API key, obtained [here](https://www.eia.gov/opendata/register.php).
|
|
51
|
+
|
|
52
|
+
### OpenBB Hub
|
|
53
|
+
|
|
54
|
+
Add the key as "eia_api_key" in the OpenBB Hub Credentials page, [here](https://my.openbb.co/app/platform/credentials)
|
|
55
|
+
|
|
56
|
+
### `user_settings.json`
|
|
57
|
+
|
|
58
|
+
Add it to the credentials section of `~/.openbb_platform/user_settings.json`
|
|
59
|
+
|
|
60
|
+
```json
|
|
61
|
+
{
|
|
62
|
+
"credentials": {
|
|
63
|
+
"eia_api_key": "REPLACE_WITH_YOUR_KEY"
|
|
64
|
+
}
|
|
65
|
+
}
|
|
66
|
+
```
|
|
67
|
+
|
|
68
|
+
### Current Python Session
|
|
69
|
+
|
|
70
|
+
The credential can be added for the current session only, after importing the OpenBB package.
|
|
71
|
+
|
|
72
|
+
```python
|
|
73
|
+
from openbb import obb
|
|
74
|
+
|
|
75
|
+
obb.user.credentials.eia_api_key = "REPLACE_WITH_YOUR_KEY"
|
|
76
|
+
```
|
|
77
|
+
|
|
78
|
+
## Coverage
|
|
79
|
+
|
|
80
|
+
### Endpoints
|
|
81
|
+
|
|
82
|
+
- `obb.commodity.petroleum_status_report` (API key not required.)
|
|
83
|
+
- `obb.commodity.short_term_energy_outlook` (API key required.)
|
|
84
|
+
|
|
85
|
+
### Weekly Petroluem Status Report
|
|
86
|
+
|
|
87
|
+
The WPSR is comprised of thirteen (excludes discontinued series) high-level categories with each containing a subset of tables. Data is from the static Excel files published [here](https://www.eia.gov/petroleum/supply/weekly/), and each file represents a single category.
|
|
88
|
+
|
|
89
|
+
All data from a single category is returned by supplying "all" to the `table` parameter of the WPSR endpoint.
|
|
90
|
+
|
|
91
|
+
Tables from the WPSR are returned in a flat format in the same order as presented in the Excel files. The response is suitable for pivot tables and SQL storage.
|
|
92
|
+
|
|
93
|
+
Category choices are defined as:
|
|
94
|
+
|
|
95
|
+
balance_sheet
|
|
96
|
+
inputs_and_production
|
|
97
|
+
refiner_and_blender_net_production
|
|
98
|
+
crude_petroleum_stocks
|
|
99
|
+
gasoline_fuel_stocks
|
|
100
|
+
total_gasoline_by_sub_padd
|
|
101
|
+
distillate_fuel_oil_stocks
|
|
102
|
+
imports
|
|
103
|
+
imports_by_country
|
|
104
|
+
weekly_estimates
|
|
105
|
+
spot_prices_crude_gas_heating
|
|
106
|
+
spot_prices_diesel_jet_fuel_propane
|
|
107
|
+
retail_prices
|
|
108
|
+
|
|
109
|
+
### Short Term Energy Outlook
|
|
110
|
+
|
|
111
|
+
The Short Term Energy Outlook (STEO) is curated by table, and relies on the EIA V2 API. Tables are defined by their alphanumeric code, and return in the same format as the WPSR tables.
|
|
112
|
+
|
|
113
|
+
01: US Energy Markets Summary
|
|
114
|
+
02: Nominal Energy Prices
|
|
115
|
+
03a: World Petroleum and Other Liquid Fuels Production, Consumption, and Inventories
|
|
116
|
+
03b: Non-OPEC Petroleum and Other Liquid Fuels Production
|
|
117
|
+
03c: World Petroleum and Other Liquid Fuels Production
|
|
118
|
+
03d: World Crude Oil Production
|
|
119
|
+
03e: World Petroleum and Other Liquid Fuels Consumption
|
|
120
|
+
04a: US Petroleum and Other Liquid Fuels Supply, Consumption, and Inventories
|
|
121
|
+
04b: US Hydrocarbon Gas Liquids (HGL) and Petroleum Refinery Balances
|
|
122
|
+
04c: US Regional Motor Gasoline Prices and Inventories
|
|
123
|
+
04d: US Biofuel Supply, Consumption, and Inventories
|
|
124
|
+
05a: US Natural Gas Supply, Consumption, and Inventories
|
|
125
|
+
05b: US Regional Natural Gas Prices
|
|
126
|
+
06: US Coal Supply, Consumption, and Inventories
|
|
127
|
+
07a: US Electricity Industry Overview
|
|
128
|
+
07b: US Regional Electricity Retail Sales
|
|
129
|
+
07c: US Regional Electricity Prices
|
|
130
|
+
07d1: US Regional Electricity Generation, Electric Power Sector
|
|
131
|
+
07d2: US Regional Electricity Generation, Electric Power Sector, continued
|
|
132
|
+
07e: US Electricity Generating Capacity
|
|
133
|
+
08: US Renewable Energy Consumption
|
|
134
|
+
09a: US Macroeconomic Indicators and CO2 Emissions
|
|
135
|
+
09b: US Regional Macroeconomic Data
|
|
136
|
+
09c: US Regional Weather Data
|
|
137
|
+
10a: Drilling Productivity Metrics
|
|
138
|
+
10b: Crude Oil and Natural Gas Production from Shale and Tight Formations
|
|
139
|
+
|
|
140
|
+
A "symbol" parameter allows lookup by individual series ID(s) within the dataset.
|
|
@@ -0,0 +1,119 @@
|
|
|
1
|
+
# OpenBB EIA Provider Extension
|
|
2
|
+
|
|
3
|
+
This module integrates the [EIA](https://eia.gov) data provider into the OpenBB Platform.
|
|
4
|
+
|
|
5
|
+
## Installation
|
|
6
|
+
|
|
7
|
+
### PyPI
|
|
8
|
+
|
|
9
|
+
```sh
|
|
10
|
+
pip install openbb-us-eia
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
### From Source
|
|
14
|
+
|
|
15
|
+
After cloning the main repository, navigate into this folder and enter:
|
|
16
|
+
|
|
17
|
+
```sh
|
|
18
|
+
pip install .
|
|
19
|
+
```
|
|
20
|
+
|
|
21
|
+
To install in editable mode:
|
|
22
|
+
|
|
23
|
+
```sh
|
|
24
|
+
pip install -e .
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
## Authorization
|
|
28
|
+
|
|
29
|
+
Functions calling the EIA's API require free registration and an API key, obtained [here](https://www.eia.gov/opendata/register.php).
|
|
30
|
+
|
|
31
|
+
### OpenBB Hub
|
|
32
|
+
|
|
33
|
+
Add the key as "eia_api_key" in the OpenBB Hub Credentials page, [here](https://my.openbb.co/app/platform/credentials)
|
|
34
|
+
|
|
35
|
+
### `user_settings.json`
|
|
36
|
+
|
|
37
|
+
Add it to the credentials section of `~/.openbb_platform/user_settings.json`
|
|
38
|
+
|
|
39
|
+
```json
|
|
40
|
+
{
|
|
41
|
+
"credentials": {
|
|
42
|
+
"eia_api_key": "REPLACE_WITH_YOUR_KEY"
|
|
43
|
+
}
|
|
44
|
+
}
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
### Current Python Session
|
|
48
|
+
|
|
49
|
+
The credential can be added for the current session only, after importing the OpenBB package.
|
|
50
|
+
|
|
51
|
+
```python
|
|
52
|
+
from openbb import obb
|
|
53
|
+
|
|
54
|
+
obb.user.credentials.eia_api_key = "REPLACE_WITH_YOUR_KEY"
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
## Coverage
|
|
58
|
+
|
|
59
|
+
### Endpoints
|
|
60
|
+
|
|
61
|
+
- `obb.commodity.petroleum_status_report` (API key not required.)
|
|
62
|
+
- `obb.commodity.short_term_energy_outlook` (API key required.)
|
|
63
|
+
|
|
64
|
+
### Weekly Petroluem Status Report
|
|
65
|
+
|
|
66
|
+
The WPSR is comprised of thirteen (excludes discontinued series) high-level categories with each containing a subset of tables. Data is from the static Excel files published [here](https://www.eia.gov/petroleum/supply/weekly/), and each file represents a single category.
|
|
67
|
+
|
|
68
|
+
All data from a single category is returned by supplying "all" to the `table` parameter of the WPSR endpoint.
|
|
69
|
+
|
|
70
|
+
Tables from the WPSR are returned in a flat format in the same order as presented in the Excel files. The response is suitable for pivot tables and SQL storage.
|
|
71
|
+
|
|
72
|
+
Category choices are defined as:
|
|
73
|
+
|
|
74
|
+
balance_sheet
|
|
75
|
+
inputs_and_production
|
|
76
|
+
refiner_and_blender_net_production
|
|
77
|
+
crude_petroleum_stocks
|
|
78
|
+
gasoline_fuel_stocks
|
|
79
|
+
total_gasoline_by_sub_padd
|
|
80
|
+
distillate_fuel_oil_stocks
|
|
81
|
+
imports
|
|
82
|
+
imports_by_country
|
|
83
|
+
weekly_estimates
|
|
84
|
+
spot_prices_crude_gas_heating
|
|
85
|
+
spot_prices_diesel_jet_fuel_propane
|
|
86
|
+
retail_prices
|
|
87
|
+
|
|
88
|
+
### Short Term Energy Outlook
|
|
89
|
+
|
|
90
|
+
The Short Term Energy Outlook (STEO) is curated by table, and relies on the EIA V2 API. Tables are defined by their alphanumeric code, and return in the same format as the WPSR tables.
|
|
91
|
+
|
|
92
|
+
01: US Energy Markets Summary
|
|
93
|
+
02: Nominal Energy Prices
|
|
94
|
+
03a: World Petroleum and Other Liquid Fuels Production, Consumption, and Inventories
|
|
95
|
+
03b: Non-OPEC Petroleum and Other Liquid Fuels Production
|
|
96
|
+
03c: World Petroleum and Other Liquid Fuels Production
|
|
97
|
+
03d: World Crude Oil Production
|
|
98
|
+
03e: World Petroleum and Other Liquid Fuels Consumption
|
|
99
|
+
04a: US Petroleum and Other Liquid Fuels Supply, Consumption, and Inventories
|
|
100
|
+
04b: US Hydrocarbon Gas Liquids (HGL) and Petroleum Refinery Balances
|
|
101
|
+
04c: US Regional Motor Gasoline Prices and Inventories
|
|
102
|
+
04d: US Biofuel Supply, Consumption, and Inventories
|
|
103
|
+
05a: US Natural Gas Supply, Consumption, and Inventories
|
|
104
|
+
05b: US Regional Natural Gas Prices
|
|
105
|
+
06: US Coal Supply, Consumption, and Inventories
|
|
106
|
+
07a: US Electricity Industry Overview
|
|
107
|
+
07b: US Regional Electricity Retail Sales
|
|
108
|
+
07c: US Regional Electricity Prices
|
|
109
|
+
07d1: US Regional Electricity Generation, Electric Power Sector
|
|
110
|
+
07d2: US Regional Electricity Generation, Electric Power Sector, continued
|
|
111
|
+
07e: US Electricity Generating Capacity
|
|
112
|
+
08: US Renewable Energy Consumption
|
|
113
|
+
09a: US Macroeconomic Indicators and CO2 Emissions
|
|
114
|
+
09b: US Regional Macroeconomic Data
|
|
115
|
+
09c: US Regional Weather Data
|
|
116
|
+
10a: Drilling Productivity Metrics
|
|
117
|
+
10b: Crude Oil and Natural Gas Production from Shale and Tight Formations
|
|
118
|
+
|
|
119
|
+
A "symbol" parameter allows lookup by individual series ID(s) within the dataset.
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
"""OpenBB EIA Provider Module."""
|
|
2
|
+
|
|
3
|
+
from openbb_core.provider.abstract.provider import Provider
|
|
4
|
+
from openbb_us_eia.models.petroleum_status_report import EiaPetroleumStatusReportFetcher
|
|
5
|
+
from openbb_us_eia.models.short_term_energy_outlook import (
|
|
6
|
+
EiaShortTermEnergyOutlookFetcher,
|
|
7
|
+
)
|
|
8
|
+
|
|
9
|
+
eia_provider = Provider(
|
|
10
|
+
name="eia",
|
|
11
|
+
website="https://eia.gov/",
|
|
12
|
+
description="The U.S. Energy Information Administration is committed to its free and open data"
|
|
13
|
+
+ " by making it available through an Application Programming Interface (API) and its open data tools."
|
|
14
|
+
+ " See https://www.eia.gov/opendata/ for more information.",
|
|
15
|
+
credentials=[
|
|
16
|
+
"api_key"
|
|
17
|
+
], # This is not required for the Weekly Petroleum Status Report
|
|
18
|
+
fetcher_dict={
|
|
19
|
+
"PetroleumStatusReport": EiaPetroleumStatusReportFetcher,
|
|
20
|
+
"ShortTermEnergyOutlook": EiaShortTermEnergyOutlookFetcher,
|
|
21
|
+
},
|
|
22
|
+
repr_name="U.S. Energy Information Administration (EIA) Open Data and API",
|
|
23
|
+
instructions="""Credentials are required for functions calling the EIA's API.
|
|
24
|
+
Register for a free key here: https://www.eia.gov/opendata/register.php""",
|
|
25
|
+
)
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""OpenBB EIA Provider Models."""
|
|
@@ -0,0 +1,241 @@
|
|
|
1
|
+
"""EIA Weekly Petroleum Status Report model."""
|
|
2
|
+
|
|
3
|
+
# pylint: disable=unused-argument
|
|
4
|
+
|
|
5
|
+
from typing import Any, Optional
|
|
6
|
+
|
|
7
|
+
from openbb_core.app.model.abstract.error import OpenBBError
|
|
8
|
+
from openbb_core.provider.abstract.fetcher import Fetcher
|
|
9
|
+
from openbb_core.provider.standard_models.petroleum_status_report import (
|
|
10
|
+
PetroleumStatusReportData,
|
|
11
|
+
PetroleumStatusReportQueryParams,
|
|
12
|
+
)
|
|
13
|
+
from openbb_core.provider.utils.errors import EmptyDataError
|
|
14
|
+
from openbb_us_eia.utils.constants import (
|
|
15
|
+
WpsrCategoryChoices,
|
|
16
|
+
WpsrCategoryType,
|
|
17
|
+
WpsrFileMap,
|
|
18
|
+
WpsrTableChoices,
|
|
19
|
+
WpsrTableMap,
|
|
20
|
+
)
|
|
21
|
+
from pydantic import Field
|
|
22
|
+
|
|
23
|
+
WpsrTableChoicesString = "\n ".join(WpsrTableChoices)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class EiaPetroleumStatusReportQueryParams(PetroleumStatusReportQueryParams):
|
|
27
|
+
"""EIA Petroleum Status Report Query Parameters.
|
|
28
|
+
|
|
29
|
+
Source: https://www.eia.gov/petroleum/supply/weekly/
|
|
30
|
+
"""
|
|
31
|
+
|
|
32
|
+
__json_schema_extra__ = {
|
|
33
|
+
"category": {
|
|
34
|
+
"multiiple_items_allowed": False,
|
|
35
|
+
"choices": WpsrCategoryChoices,
|
|
36
|
+
},
|
|
37
|
+
"table": {
|
|
38
|
+
"multiple_items_allowed": True,
|
|
39
|
+
"choices": WpsrTableChoices,
|
|
40
|
+
},
|
|
41
|
+
}
|
|
42
|
+
|
|
43
|
+
category: WpsrCategoryType = Field(
|
|
44
|
+
default="balance_sheet",
|
|
45
|
+
description="The group of data to be returned. The default is the balance sheet.",
|
|
46
|
+
)
|
|
47
|
+
table: Optional[str] = Field(
|
|
48
|
+
default=None,
|
|
49
|
+
description="The specific table element within the category to be returned,"
|
|
50
|
+
+ " default is 'stocks', if the category is 'weekly_estimates', else 'all'."
|
|
51
|
+
+ "\n Note: Choices represent all available tables from the entire collection and are not all"
|
|
52
|
+
+ " available for every category."
|
|
53
|
+
+ "\n Invalid choices will raise a ValidationError with a message"
|
|
54
|
+
+ " indicating the valid choices for the selected category."
|
|
55
|
+
+ "\n Choices are:"
|
|
56
|
+
+ f"\n {WpsrTableChoicesString}\n ",
|
|
57
|
+
)
|
|
58
|
+
use_cache: bool = Field(
|
|
59
|
+
default=True,
|
|
60
|
+
description="Subsequent requests for the same source data are cached for the session using ALRU cache.",
|
|
61
|
+
)
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
class EiaPetroleumStatusReportData(PetroleumStatusReportData):
|
|
65
|
+
"""EIA Petroleum Status Report Data Model."""
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
class EiaPetroleumStatusReportFetcher(
|
|
69
|
+
Fetcher[EiaPetroleumStatusReportQueryParams, list[EiaPetroleumStatusReportData]]
|
|
70
|
+
):
|
|
71
|
+
"""EIA Petroleum Status Report Fetcher."""
|
|
72
|
+
|
|
73
|
+
require_credentials = False
|
|
74
|
+
|
|
75
|
+
@staticmethod
|
|
76
|
+
def transform_query(params: dict[str, Any]) -> EiaPetroleumStatusReportQueryParams:
|
|
77
|
+
"""Transform the query parameters."""
|
|
78
|
+
# pylint: disable=import-outside-toplevel
|
|
79
|
+
from warnings import warn
|
|
80
|
+
|
|
81
|
+
category = params.get("category", "balance_sheet")
|
|
82
|
+
tables = WpsrTableMap.get(category, {})
|
|
83
|
+
_table = params.get("table", "")
|
|
84
|
+
|
|
85
|
+
if not _table:
|
|
86
|
+
_table = "stocks" if category == "weekly_estimates" else "all"
|
|
87
|
+
|
|
88
|
+
_tables = _table.split(",")
|
|
89
|
+
|
|
90
|
+
if len(_tables) == 1 and _tables[0] == "all" and category == "weekly_estimates":
|
|
91
|
+
raise OpenBBError(
|
|
92
|
+
ValueError(
|
|
93
|
+
f"'all' is not a supported choice for {category}. Please choose from: {list(tables)}"
|
|
94
|
+
)
|
|
95
|
+
)
|
|
96
|
+
|
|
97
|
+
if "all" in _tables and len(_tables) > 1:
|
|
98
|
+
_tables.remove("all")
|
|
99
|
+
warn("'all' cannot be used with other table choices. Ignoring 'all'.")
|
|
100
|
+
|
|
101
|
+
for table in _tables:
|
|
102
|
+
if table != "all" and table not in tables:
|
|
103
|
+
raise OpenBBError(
|
|
104
|
+
ValueError(
|
|
105
|
+
f"Invalid table choice: {table}. Valid choices for {category}: {list(tables)}"
|
|
106
|
+
)
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
params["table"] = ",".join(_tables)
|
|
110
|
+
|
|
111
|
+
return EiaPetroleumStatusReportQueryParams(**params)
|
|
112
|
+
|
|
113
|
+
@staticmethod
|
|
114
|
+
async def aextract_data(
|
|
115
|
+
query: EiaPetroleumStatusReportQueryParams,
|
|
116
|
+
credentials: Optional[dict[str, Any]],
|
|
117
|
+
**kwargs: Any,
|
|
118
|
+
) -> dict:
|
|
119
|
+
"""Extract the data from the EIA website."""
|
|
120
|
+
# pylint: disable=import-outside-toplevel
|
|
121
|
+
from openbb_us_eia.utils.helpers import download_excel_file
|
|
122
|
+
|
|
123
|
+
url = WpsrFileMap.get(query.category, "balance_sheet")
|
|
124
|
+
|
|
125
|
+
try:
|
|
126
|
+
results = await download_excel_file(url, query.use_cache)
|
|
127
|
+
except OpenBBError as e:
|
|
128
|
+
raise OpenBBError(f"Error extracting data -> {e}") from e
|
|
129
|
+
|
|
130
|
+
return {"file": results}
|
|
131
|
+
|
|
132
|
+
@staticmethod
|
|
133
|
+
def transform_data(
|
|
134
|
+
query: EiaPetroleumStatusReportQueryParams,
|
|
135
|
+
data: dict,
|
|
136
|
+
**kwargs: Any,
|
|
137
|
+
) -> list[EiaPetroleumStatusReportData]:
|
|
138
|
+
"""Transform the data."""
|
|
139
|
+
# pylint: disable=import-outside-toplevel
|
|
140
|
+
import concurrent.futures # noqa
|
|
141
|
+
import re
|
|
142
|
+
from functools import lru_cache
|
|
143
|
+
from numpy import nan
|
|
144
|
+
from pandas import Categorical, ExcelFile, concat, read_excel
|
|
145
|
+
from warnings import warn
|
|
146
|
+
|
|
147
|
+
category = query.category
|
|
148
|
+
|
|
149
|
+
_tables = (
|
|
150
|
+
query.table.split(",") # type: ignore
|
|
151
|
+
if query.table
|
|
152
|
+
else ["stocks"] if category == "weekly_estimates" else ["all"]
|
|
153
|
+
)
|
|
154
|
+
all_tables = list(WpsrTableMap[category])
|
|
155
|
+
tables = all_tables if "all" in _tables else _tables
|
|
156
|
+
|
|
157
|
+
file = data.get("file")
|
|
158
|
+
|
|
159
|
+
if not isinstance(file, ExcelFile):
|
|
160
|
+
raise OpenBBError(
|
|
161
|
+
TypeError(f"Expected an ExcelFile object, got {type(file)} instead.")
|
|
162
|
+
)
|
|
163
|
+
|
|
164
|
+
dfs: list = []
|
|
165
|
+
|
|
166
|
+
def replace_data_strings(text):
|
|
167
|
+
"""Replace the table strings with sortable numbers."""
|
|
168
|
+
pattern = r"Data (\d):"
|
|
169
|
+
|
|
170
|
+
def replacer(match):
|
|
171
|
+
"""Replace the matched string with a sortable number."""
|
|
172
|
+
return f"Data 0{match.group(1)}:"
|
|
173
|
+
|
|
174
|
+
return re.sub(pattern, replacer, text)
|
|
175
|
+
|
|
176
|
+
@lru_cache(maxsize=128)
|
|
177
|
+
def read_excel_file(file, category, table):
|
|
178
|
+
"""Read the ExcelFile for the sheet name and flatten the table."""
|
|
179
|
+
sheet_name = WpsrTableMap[category][table]
|
|
180
|
+
table_name = read_excel(file, sheet_name, header=None, nrows=1).iloc[0, 1]
|
|
181
|
+
table_name = replace_data_strings(table_name)
|
|
182
|
+
df = read_excel(file, sheet_name, header=[1, 2], nrows=3)
|
|
183
|
+
symbols = df.columns.get_level_values(0).tolist()
|
|
184
|
+
titles = [
|
|
185
|
+
d.replace(".1", "") for d in df.columns.get_level_values(1).tolist()
|
|
186
|
+
]
|
|
187
|
+
title_map = dict(zip(symbols, titles))
|
|
188
|
+
df = read_excel(file, sheet_name, header=None, skiprows=3)
|
|
189
|
+
df.columns = [d.replace("Sourcekey", "date") for d in symbols]
|
|
190
|
+
df = df.melt(
|
|
191
|
+
id_vars="date",
|
|
192
|
+
value_vars=[d for d in df.columns if d != "date"],
|
|
193
|
+
var_name="symbol",
|
|
194
|
+
).dropna()
|
|
195
|
+
df = df.reset_index(drop=True)
|
|
196
|
+
df.loc[:, "title"] = df.symbol.map(title_map)
|
|
197
|
+
df.loc[:, "unit"] = df.title.map(lambda x: x.split(" (")[-1].split(")")[0])
|
|
198
|
+
units = [f"({d})" for d in df.unit.unique().tolist()]
|
|
199
|
+
for unit in units:
|
|
200
|
+
df.title = df.title.str.replace(unit, "", regex=False).str.strip()
|
|
201
|
+
df.loc[:, "table"] = table_name
|
|
202
|
+
df["order"] = df.groupby("date").cumcount() + 1
|
|
203
|
+
df = df[["date", "table", "symbol", "order", "title", "value", "unit"]]
|
|
204
|
+
df.symbol = Categorical(df.symbol, categories=symbols, ordered=True)
|
|
205
|
+
df = df.sort_values(["date", "symbol"])
|
|
206
|
+
df.date = df.date.dt.date
|
|
207
|
+
|
|
208
|
+
if query.start_date:
|
|
209
|
+
df = df[df.date >= query.start_date]
|
|
210
|
+
|
|
211
|
+
if query.end_date:
|
|
212
|
+
df = df[df.date <= query.end_date]
|
|
213
|
+
|
|
214
|
+
df = df.reset_index(drop=True)
|
|
215
|
+
|
|
216
|
+
if len(df) > 0:
|
|
217
|
+
dfs.append(df)
|
|
218
|
+
else:
|
|
219
|
+
warn(f"No data for table: {table}")
|
|
220
|
+
|
|
221
|
+
try:
|
|
222
|
+
with concurrent.futures.ThreadPoolExecutor() as executor:
|
|
223
|
+
executor.map(
|
|
224
|
+
lambda table: read_excel_file(file, category, table), tables
|
|
225
|
+
)
|
|
226
|
+
|
|
227
|
+
results = concat(dfs)
|
|
228
|
+
|
|
229
|
+
if len(results) < 1:
|
|
230
|
+
raise EmptyDataError("The data is empty.")
|
|
231
|
+
|
|
232
|
+
results = results.sort_values(by=["date", "table", "order"]).replace(
|
|
233
|
+
{nan: None}
|
|
234
|
+
)
|
|
235
|
+
|
|
236
|
+
return [
|
|
237
|
+
EiaPetroleumStatusReportData.model_validate(d)
|
|
238
|
+
for d in results.to_dict(orient="records")
|
|
239
|
+
]
|
|
240
|
+
except Exception as e: # pylint: disable=broad-except
|
|
241
|
+
raise OpenBBError(f"Error transforming the data -> {e}") from e
|