Skip to content

Commit fab4b5c

Browse files
ahuang11ahuang11
committed
indentation + ruff
1 parent b9a1e26 commit fab4b5c

15 files changed

Lines changed: 2760 additions & 2724 deletions

‎perf_tests/compute_air.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,10 @@
44
import xarray_sql as xql
55

66
if __name__ == "__main__":
7-
air = xr.tutorial.open_dataset("air_temperature")
8-
chunks = {"time": 240}
9-
air = air.chunk(chunks)
7+
air = xr.tutorial.open_dataset("air_temperature")
8+
chunks = {"time": 240}
9+
air = air.chunk(chunks)
1010

11-
df = xql.read_xarray(air).read_pandas()
11+
df = xql.read_xarray(air).read_pandas()
1212

13-
print(len(df))
13+
print(len(df))

‎perf_tests/groupby_air.py‎

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -6,32 +6,32 @@
66

77

88
if __name__ == "__main__":
9-
air = xr.tutorial.open_dataset("air_temperature")
10-
chunks = {"time": 240, "lat": 5, "lon": 7}
11-
air = air.chunk(chunks)
12-
air_small = air.isel(
13-
time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)
14-
).chunk(chunks)
9+
air = xr.tutorial.open_dataset("air_temperature")
10+
chunks = {"time": 240, "lat": 5, "lon": 7}
11+
air = air.chunk(chunks)
12+
air_small = air.isel(
13+
time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)
14+
).chunk(chunks)
1515

16-
df = xql.read_xarray_table(air_small)
16+
df = xql.read_xarray_table(air_small)
1717

18-
ctx = SessionContext()
19-
ctx.register_table("air", df)
18+
ctx = SessionContext()
19+
ctx.register_table("air", df)
2020

21-
query = ctx.sql(
22-
"""
21+
query = ctx.sql(
22+
"""
2323
SELECT
2424
"lat", "lon", SUM("air") as air_total
2525
FROM
2626
"air"
2727
GROUP BY
2828
"lat", "lon"
2929
"""
30-
)
30+
)
3131

32-
result = query.collect()
32+
result = query.collect()
3333

34-
expected = air_small.sizes["lat"] * air_small.sizes["lon"]
35-
actual = sum(len(batch) for batch in result)
36-
assert actual == expected, f"Length must be {expected}, but was {actual}."
37-
print(expected)
34+
expected = air_small.sizes["lat"] * air_small.sizes["lon"]
35+
actual = sum(len(batch) for batch in result)
36+
assert actual == expected, f"Length must be {expected}, but was {actual}."
37+
print(expected)

‎perf_tests/groupby_air_full.py‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -5,30 +5,30 @@
55
from datafusion import SessionContext
66

77
if __name__ == "__main__":
8-
air = xr.tutorial.open_dataset("air_temperature")
9-
chunks = {"time": 240}
10-
air = air.chunk(chunks)
8+
air = xr.tutorial.open_dataset("air_temperature")
9+
chunks = {"time": 240}
10+
air = air.chunk(chunks)
1111

12-
df = xql.read_xarray_table(air)
12+
df = xql.read_xarray_table(air)
1313

14-
ctx = SessionContext()
15-
ctx.register_table("air", df)
14+
ctx = SessionContext()
15+
ctx.register_table("air", df)
1616

17-
query = ctx.sql(
18-
"""
17+
query = ctx.sql(
18+
"""
1919
SELECT
2020
"lat", "lon", SUM("air") as air_total
2121
FROM
2222
"air"
2323
GROUP BY
2424
"lat", "lon"
2525
"""
26-
)
26+
)
2727

28-
result = query.collect()
28+
result = query.collect()
2929

30-
expected = air.sizes["lat"] * air.sizes["lon"]
31-
actual = sum(len(batch) for batch in result)
30+
expected = air.sizes["lat"] * air.sizes["lon"]
31+
actual = sum(len(batch) for batch in result)
3232

33-
assert actual == expected, f"Length must be {expected}, but was {actual}."
34-
print(expected)
33+
assert actual == expected, f"Length must be {expected}, but was {actual}."
34+
print(expected)

‎perf_tests/sanity.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,13 @@
44
import xarray_sql as xql
55

66
if __name__ == "__main__":
7-
air = xr.tutorial.open_dataset("air_temperature")
8-
chunks = {"time": 240, "lat": 5, "lon": 7}
7+
air = xr.tutorial.open_dataset("air_temperature")
8+
chunks = {"time": 240, "lat": 5, "lon": 7}
99

10-
air_small = air.isel(
11-
time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)
12-
).chunk(chunks)
10+
air_small = air.isel(
11+
time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)
12+
).chunk(chunks)
1313

14-
df = xql.read_xarray(air_small).read_pandas()
14+
df = xql.read_xarray(air_small).read_pandas()
1515

16-
print(len(df))
16+
print(len(df))

‎pyproject.toml‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -64,11 +64,13 @@ module-name = "xarray_sql._native"
6464
[tool.setuptools.packages.find]
6565
exclude = ["demo", "perf_tests", "tests", "tests.*"]
6666

67-
[tool.pyink]
67+
[tool.ruff]
6868
line-length = 80
69-
preview = true
70-
pyink-indentation = 2
71-
pyink-use-majority-quotes = true
69+
indent-width = 4
70+
71+
[tool.ruff.format]
72+
indent-style = "space"
73+
quote-style = "double"
7274

7375
[tool.mypy]
7476
python_version = "3.11"
@@ -98,7 +100,7 @@ dev = [
98100
"xarray_sql[test]",
99101
"xarray_sql[docs]",
100102
"py-spy>=0.4.0",
101-
"pyink>=24.10.1",
103+
"ruff>=0.15.10",
102104
"maturin>=1.9.1",
103105
]
104106

‎tests/conftest.py‎

Lines changed: 92 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -6,139 +6,145 @@
66

77

88
def rand_wx(start: str, end: str) -> xr.Dataset:
9-
np.random.seed(42)
10-
lat = np.linspace(-90, 90, num=720)
11-
lon = np.linspace(-180, 180, num=1440)
12-
time = pd.date_range(start, end, freq="h")
13-
level = np.array([1000, 500], dtype=np.int32)
14-
reference_time = pd.Timestamp(start)
15-
temperature = 15 + 8 * np.random.randn(720, 1440, len(time), len(level))
16-
precipitation = 10 * np.random.rand(720, 1440, len(time), len(level))
17-
return xr.Dataset(
18-
data_vars=dict(
19-
temperature=(["lat", "lon", "time", "level"], temperature),
20-
precipitation=(["lat", "lon", "time", "level"], precipitation),
21-
),
22-
coords=dict(
23-
lat=lat,
24-
lon=lon,
25-
time=time,
26-
level=level,
27-
reference_time=reference_time,
28-
),
29-
attrs=dict(description="Random weather."),
30-
)
9+
np.random.seed(42)
10+
lat = np.linspace(-90, 90, num=720)
11+
lon = np.linspace(-180, 180, num=1440)
12+
time = pd.date_range(start, end, freq="h")
13+
level = np.array([1000, 500], dtype=np.int32)
14+
reference_time = pd.Timestamp(start)
15+
temperature = 15 + 8 * np.random.randn(720, 1440, len(time), len(level))
16+
precipitation = 10 * np.random.rand(720, 1440, len(time), len(level))
17+
return xr.Dataset(
18+
data_vars=dict(
19+
temperature=(["lat", "lon", "time", "level"], temperature),
20+
precipitation=(["lat", "lon", "time", "level"], precipitation),
21+
),
22+
coords=dict(
23+
lat=lat,
24+
lon=lon,
25+
time=time,
26+
level=level,
27+
reference_time=reference_time,
28+
),
29+
attrs=dict(description="Random weather."),
30+
)
3131

3232

3333
def create_large_dataset(time_steps=1000, lat_points=100, lon_points=100):
34-
"""Create a large xarray dataset for memory testing."""
35-
np.random.seed(42)
34+
"""Create a large xarray dataset for memory testing."""
35+
np.random.seed(42)
3636

37-
time = pd.date_range("2020-01-01", periods=time_steps, freq="h")
38-
lat = np.linspace(-90, 90, lat_points)
39-
lon = np.linspace(-180, 180, lon_points)
37+
time = pd.date_range("2020-01-01", periods=time_steps, freq="h")
38+
lat = np.linspace(-90, 90, lat_points)
39+
lon = np.linspace(-180, 180, lon_points)
4040

41-
temp_data = np.random.rand(time_steps, lat_points, lon_points) * 40 - 10
42-
precip_data = np.random.rand(time_steps, lat_points, lon_points) * 100
41+
temp_data = np.random.rand(time_steps, lat_points, lon_points) * 40 - 10
42+
precip_data = np.random.rand(time_steps, lat_points, lon_points) * 100
4343

44-
return xr.Dataset(
45-
{
46-
"temperature": (["time", "lat", "lon"], temp_data),
47-
"precipitation": (["time", "lat", "lon"], precip_data),
48-
},
49-
coords={"time": time, "lat": lat, "lon": lon},
50-
)
44+
return xr.Dataset(
45+
{
46+
"temperature": (["time", "lat", "lon"], temp_data),
47+
"precipitation": (["time", "lat", "lon"], precip_data),
48+
},
49+
coords={"time": time, "lat": lat, "lon": lon},
50+
)
5151

5252

5353
@pytest.fixture
5454
def air():
55-
ds = xr.tutorial.open_dataset("air_temperature")
56-
chunks = {"time": 240}
57-
return ds.chunk(chunks)
55+
ds = xr.tutorial.open_dataset("air_temperature")
56+
chunks = {"time": 240}
57+
return ds.chunk(chunks)
5858

5959

6060
@pytest.fixture
6161
def air_small(air):
62-
return air.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)).chunk(
63-
{"time": 240}
64-
)
62+
return air.isel(
63+
time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)
64+
).chunk({"time": 240})
6565

6666

6767
@pytest.fixture
6868
def randwx():
69-
return rand_wx("1995-01-13T00", "1995-01-13T01")
69+
return rand_wx("1995-01-13T00", "1995-01-13T01")
7070

7171

7272
@pytest.fixture
7373
def large_ds():
74-
return create_large_dataset().chunk({"time": 25})
74+
return create_large_dataset().chunk({"time": 25})
7575

7676

7777
@pytest.fixture
7878
def air_dataset_small():
79-
ds = xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
80-
return ds.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10))
79+
ds = xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
80+
return ds.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10))
8181

8282

8383
@pytest.fixture
8484
def air_dataset_large():
85-
return xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
85+
return xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
8686

8787

8888
@pytest.fixture
8989
def rasm_ds():
90-
"""rasm uses cftime.DatetimeNoLeap (noleap / 365_day) for time."""
91-
return xr.tutorial.open_dataset("rasm")
90+
"""rasm uses cftime.DatetimeNoLeap (noleap / 365_day) for time."""
91+
return xr.tutorial.open_dataset("rasm")
9292

9393

9494
@pytest.fixture
9595
def weather_dataset():
96-
ds = rand_wx("2023-01-01T00", "2023-01-01T12")
97-
return ds.isel(time=slice(0, 6), lat=slice(0, 10), lon=slice(0, 10)).chunk(
98-
{"time": 3}
99-
)
96+
ds = rand_wx("2023-01-01T00", "2023-01-01T12")
97+
return ds.isel(time=slice(0, 6), lat=slice(0, 10), lon=slice(0, 10)).chunk(
98+
{"time": 3}
99+
)
100100

101101

102102
@pytest.fixture
103103
def synthetic_dataset():
104-
return create_large_dataset(
105-
time_steps=50, lat_points=20, lon_points=20
106-
).chunk({"time": 25})
104+
return create_large_dataset(
105+
time_steps=50, lat_points=20, lon_points=20
106+
).chunk({"time": 25})
107107

108108

109109
@pytest.fixture
110110
def station_dataset():
111-
return xr.Dataset(
112-
{
113-
"station_id": (["station"], [1, 2, 3, 4, 5]),
114-
"elevation": (["station"], [100, 250, 500, 750, 1000]),
115-
"name": (
116-
["station"],
117-
["Station_A", "Station_B", "Station_C", "Station_D", "Station_E"],
118-
),
119-
}
120-
).chunk({"station": 5})
111+
return xr.Dataset(
112+
{
113+
"station_id": (["station"], [1, 2, 3, 4, 5]),
114+
"elevation": (["station"], [100, 250, 500, 750, 1000]),
115+
"name": (
116+
["station"],
117+
[
118+
"Station_A",
119+
"Station_B",
120+
"Station_C",
121+
"Station_D",
122+
"Station_E",
123+
],
124+
),
125+
}
126+
).chunk({"station": 5})
121127

122128

123129
@pytest.fixture
124130
def air_and_stations():
125-
air = (
126-
xr.tutorial.open_dataset("air_temperature")
127-
.isel(time=slice(0, 12), lat=slice(0, 5), lon=slice(0, 8))
128-
.chunk({"time": 6})
129-
)
130-
stations = xr.Dataset(
131-
{
132-
"station_id": (["station"], [101, 102, 103]),
133-
"lat": (
134-
["station"],
135-
[air.lat.values[0], air.lat.values[2], air.lat.values[4]],
136-
),
137-
"lon": (
138-
["station"],
139-
[air.lon.values[1], air.lon.values[3], air.lon.values[5]],
140-
),
141-
"elevation": (["station"], [100, 250, 500]),
142-
}
143-
).chunk({"station": 3})
144-
return air, stations
131+
air = (
132+
xr.tutorial.open_dataset("air_temperature")
133+
.isel(time=slice(0, 12), lat=slice(0, 5), lon=slice(0, 8))
134+
.chunk({"time": 6})
135+
)
136+
stations = xr.Dataset(
137+
{
138+
"station_id": (["station"], [101, 102, 103]),
139+
"lat": (
140+
["station"],
141+
[air.lat.values[0], air.lat.values[2], air.lat.values[4]],
142+
),
143+
"lon": (
144+
["station"],
145+
[air.lon.values[1], air.lon.values[3], air.lon.values[5]],
146+
),
147+
"elevation": (["station"], [100, 250, 500]),
148+
}
149+
).chunk({"station": 3})
150+
return air, stations

0 commit comments

Comments
 (0)