@@ -14,17 +14,25 @@ def willr(
14
14
) -> Union [pd .DataFrame , pl .DataFrame ]:
15
15
16
16
if high_column not in data .columns :
17
- raise PyIndicatorException (f"Column '{ high_column } ' not found in DataFrame" )
17
+ raise PyIndicatorException (
18
+ f"Column '{ high_column } ' not found in DataFrame"
19
+ )
18
20
19
21
if low_column not in data .columns :
20
- raise PyIndicatorException (f"Column '{ low_column } ' not found in DataFrame" )
22
+ raise PyIndicatorException (
23
+ f"Column '{ low_column } ' not found in DataFrame"
24
+ )
21
25
22
26
if isinstance (data , pd .DataFrame ):
23
- data ["high_n" ] = data [high_column ].rolling (window = period , min_periods = 1 ).max ()
24
- data ["low_n" ] = data [low_column ].rolling (window = period , min_periods = 1 ).min ()
27
+ data ["high_n" ] = data [high_column ]\
28
+ .rolling (window = period , min_periods = 1 ).max ()
29
+ data ["low_n" ] = data [low_column ]\
30
+ .rolling (window = period , min_periods = 1 ).min ()
25
31
26
- data [result_column ] = ((data ["high_n" ] - data [close_column ]) /
27
- (data ["high_n" ] - data ["low_n" ])) * - 100
32
+ data [result_column ] = (
33
+ (data ["high_n" ] - data [close_column ]) /
34
+ (data ["high_n" ] - data ["low_n" ])
35
+ ) * - 100
28
36
29
37
# Set the first `period` rows to 0 using .iloc
30
38
if not data .empty :
@@ -33,22 +41,28 @@ def willr(
33
41
return data .drop (columns = ["high_n" , "low_n" ])
34
42
35
43
elif isinstance (data , pl .DataFrame ):
36
- high_n = data .select (pl .col (high_column ).rolling_max (period ).alias ("high_n" ))
37
- low_n = data .select (pl .col (low_column ).rolling_min (period ).alias ("low_n" ))
44
+ high_n = data .select (
45
+ pl .col (high_column ).rolling_max (period ).alias ("high_n" )
46
+ )
47
+ low_n = data .select (
48
+ pl .col (low_column ).rolling_min (period ).alias ("low_n" )
49
+ )
38
50
39
51
data = data .with_columns ([
40
52
high_n ["high_n" ],
41
53
low_n ["low_n" ]
42
54
])
43
55
44
56
data = data .with_columns (
45
- ((pl .col ("high_n" ) - pl .col (close_column )) / (pl .col ("high_n" ) - pl .col ("low_n" )) * - 100 )
57
+ ((pl .col ("high_n" ) - pl .col (close_column ))
58
+ / (pl .col ("high_n" ) - pl .col ("low_n" )) * - 100 )
46
59
.alias (result_column )
47
60
)
48
61
49
62
# Set the first `period` rows of result_column to 0 directly in Polars
50
63
if data .height > 0 :
51
- zero_values = [0 ] * (period - 1 )+ data [result_column ].to_list ()[period - 1 :]
64
+ zero_values = [0 ] * (period - 1 ) \
65
+ + data [result_column ].to_list ()[period - 1 :]
52
66
data = data .with_columns (pl .Series (result_column , zero_values ))
53
67
54
68
return data .drop (["high_n" , "low_n" ])
0 commit comments