@@ -14,17 +14,25 @@ def willr(
1414) -> Union [pd .DataFrame , pl .DataFrame ]:
1515
1616 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+ )
1820
1921 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+ )
2125
2226 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 ()
2531
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
2836
2937 # Set the first `period` rows to 0 using .iloc
3038 if not data .empty :
@@ -33,22 +41,28 @@ def willr(
3341 return data .drop (columns = ["high_n" , "low_n" ])
3442
3543 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+ )
3850
3951 data = data .with_columns ([
4052 high_n ["high_n" ],
4153 low_n ["low_n" ]
4254 ])
4355
4456 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 )
4659 .alias (result_column )
4760 )
4861
4962 # Set the first `period` rows of result_column to 0 directly in Polars
5063 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 :]
5266 data = data .with_columns (pl .Series (result_column , zero_values ))
5367
5468 return data .drop (["high_n" , "low_n" ])
0 commit comments