1use crate::model::Bar;
33use crate::studies::palette;
34use crate::studies::{Indicator, IndicatorValue, IndicatorValues};
35use egui::Color32;
36use std::sync::Arc;
37
38pub type CalculateFn = Arc<dyn Fn(&[Bar]) -> Vec<IndicatorValue> + Send + Sync>;
43
44pub type InputCalculateFn =
46 Arc<dyn Fn(&[Bar], &[super::StudyInput]) -> Vec<IndicatorValue> + Send + Sync>;
47
48pub type PlotColorFn = Arc<dyn Fn(usize, usize, &IndicatorValues) -> Option<Color32> + Send + Sync>;
50
51#[derive(Clone)]
56pub struct CustomIndicator {
57 name: String,
58 desc: String,
59 calculate_fn: CalculateFn,
60 input_calculate_fn: Option<InputCalculateFn>,
61 inputs: Vec<super::StudyInput>,
62 values: IndicatorValues,
63 colors: Vec<Color32>,
64 is_overlay: bool,
65 line_count: usize,
66 line_names: Vec<String>,
67 visible: bool,
68 value_format: Option<crate::scales::StudyFormat>,
69 scale: Option<super::StudyScale>,
70 plot_types: Vec<super::PlotType>,
71 offsets: Vec<isize>,
72 fills: Vec<super::StudyFill>,
73 plot_color: Option<PlotColorFn>,
74}
75
76impl CustomIndicator {
77 pub fn new<F>(name: impl Into<String>, calculate_fn: F) -> Self
85 where
86 F: Fn(&[Bar]) -> Vec<IndicatorValue> + Send + Sync + 'static,
87 {
88 let name = name.into();
89 Self {
90 line_names: vec![name.clone()],
91 name,
92 desc: String::new(),
93 calculate_fn: Arc::new(calculate_fn),
94 input_calculate_fn: None,
95 inputs: Vec::new(),
96 values: IndicatorValues::new(),
97 colors: vec![palette::YELLOW], is_overlay: true,
99 line_count: 1,
100 visible: true,
101 value_format: None,
102 scale: None,
103 plot_types: Vec::new(),
104 offsets: Vec::new(),
105 fills: Vec::new(),
106 plot_color: None,
107 }
108 }
109
110 pub fn with_inputs<F>(
113 mut self,
114 inputs: Vec<super::StudyInput>,
115 calculate: F,
116 ) -> Result<Self, super::StudyError>
117 where
118 F: Fn(&[Bar], &[super::StudyInput]) -> Vec<IndicatorValue> + Send + Sync + 'static,
119 {
120 for (at, input) in inputs.iter().enumerate() {
121 if input.id.is_empty() {
122 return Err(super::StudyError::EmptyInputId { index: at });
123 }
124 let reason = if inputs[..at].iter().any(|other| other.id == input.id) {
125 "is declared twice".into()
126 } else if !input.accepts(&input.value) {
127 format!("has value {:?}, expected {:?}", input.value, input.kind)
128 } else if !input.accepts(&input.default) {
129 format!("has default {:?}, expected {:?}", input.default, input.kind)
130 } else {
131 continue;
132 };
133 return Err(super::StudyError::InvalidInput {
134 id: input.id.clone(),
135 reason,
136 });
137 }
138 self.inputs = inputs;
139 self.input_calculate_fn = Some(Arc::new(calculate));
140 self.values.clear();
141 Ok(self)
142 }
143
144 pub fn with_description(mut self, desc: impl Into<String>) -> Self {
146 self.desc = desc.into();
147 self
148 }
149
150 pub fn with_color(mut self, color: Color32) -> Self {
152 self.colors = vec![color];
153 self
154 }
155
156 pub fn with_colors(mut self, colors: Vec<Color32>) -> Self {
158 self.colors = colors;
159 self
160 }
161
162 pub fn with_overlay(mut self, is_overlay: bool) -> Self {
165 self.is_overlay = is_overlay;
166 self
167 }
168
169 pub fn with_line_count(mut self, count: usize) -> Self {
171 self.line_count = count;
172 self
173 }
174
175 pub fn with_value_format(mut self, format: crate::scales::StudyFormat) -> Self {
177 self.value_format = Some(format);
178 self
179 }
180
181 pub fn with_scale(mut self, scale: super::StudyScale) -> Self {
183 self.scale = Some(scale);
184 self
185 }
186
187 pub fn with_plot_types(mut self, types: Vec<super::PlotType>) -> Self {
189 self.plot_types = types;
190 self
191 }
192
193 pub fn with_offsets(mut self, offsets: Vec<isize>) -> Self {
195 self.offsets = offsets;
196 self
197 }
198
199 pub fn with_fills(mut self, fills: Vec<super::StudyFill>) -> Self {
201 self.fills = fills;
202 self
203 }
204
205 pub fn with_plot_color<F>(mut self, color: F) -> Self
207 where
208 F: Fn(usize, usize, &IndicatorValues) -> Option<Color32> + Send + Sync + 'static,
209 {
210 self.plot_color = Some(Arc::new(color));
211 self
212 }
213
214 pub fn with_line_names(mut self, names: Vec<String>) -> Self {
216 self.line_names = names;
217 self
218 }
219}
220
221impl Indicator for CustomIndicator {
222 fn inputs(&self) -> Vec<super::StudyInput> {
223 self.inputs.clone()
224 }
225
226 fn set_input(
227 &mut self,
228 id: &str,
229 value: super::InputValue,
230 ) -> Result<(), crate::studies::StudyError> {
231 let input = self
232 .inputs
233 .iter_mut()
234 .find(|input| input.id == id)
235 .ok_or_else(|| crate::studies::StudyError::UnknownInput(id.into()))?;
236 if !input.accepts(&value) {
237 return Err(crate::studies::StudyError::invalid_input(id));
238 }
239 input.value = value;
240 self.values.clear();
241 Ok(())
242 }
243
244 fn value_format(&self) -> crate::scales::StudyFormat {
245 self.value_format.unwrap_or(if self.is_overlay {
246 crate::scales::StudyFormat::Price
247 } else {
248 crate::scales::StudyFormat::Fixed(2)
249 })
250 }
251
252 fn scale(&self) -> super::StudyScale {
253 self.scale.unwrap_or(if self.is_overlay {
254 super::StudyScale::Main
255 } else {
256 super::StudyScale::Independent
257 })
258 }
259
260 fn plot_style(&self, line: usize) -> super::PlotStyle {
261 super::PlotStyle {
262 color: self.colors.get(line).copied().unwrap_or(Color32::WHITE),
263 width: 1.0,
264 pattern: crate::model::LineStyle::Solid,
265 plot_type: self.plot_types.get(line).copied().unwrap_or_default(),
266 visible: true,
267 }
268 }
269
270 fn plot_offset(&self, line: usize) -> isize {
271 self.offsets.get(line).copied().unwrap_or(0)
272 }
273
274 fn plot_color(&self, line: usize, bar: usize) -> Option<Color32> {
275 self.plot_color
276 .as_ref()
277 .and_then(|color| color(line, bar, &self.values))
278 }
279
280 fn fills(&self) -> Vec<super::StudyFill> {
281 self.fills.clone()
282 }
283
284 fn name(&self) -> &str {
285 &self.name
286 }
287
288 fn desc(&self) -> &str {
289 &self.desc
290 }
291
292 fn calculate(&mut self, data: &[Bar]) {
293 self.values = match &self.input_calculate_fn {
294 Some(calculate) => calculate(data, &self.inputs),
295 None => (self.calculate_fn)(data),
296 }
297 .into();
298 }
299
300 fn values(&self) -> &IndicatorValues {
301 &self.values
302 }
303
304 fn colors(&self) -> Vec<Color32> {
305 self.colors.clone()
306 }
307
308 fn set_colors(&mut self, colors: Vec<Color32>) {
309 if !colors.is_empty() {
310 self.colors = colors;
311 }
312 }
313
314 fn is_overlay(&self) -> bool {
315 self.is_overlay
316 }
317
318 fn line_count(&self) -> usize {
319 self.line_count
320 }
321
322 fn line_names(&self) -> Vec<String> {
323 self.line_names.clone()
324 }
325
326 fn is_visible(&self) -> bool {
327 self.visible
328 }
329
330 fn set_visible(&mut self, visible: bool) {
331 self.visible = visible;
332 }
333
334 fn clone_box(&self) -> Box<dyn Indicator> {
335 Box::new(self.clone())
336 }
337}
338
339#[cfg(test)]
340mod tests {
341 use super::*;
342
343 #[test]
344 fn declared_inputs_validate_edits_and_reach_cloned_calculations() {
345 use crate::model::{PriceSource, Timeframe};
346 use crate::studies::{InputKind, InputValue, StudyInput, StudySource};
347 let schema = vec![
348 StudyInput::new(
349 "length",
350 "Length",
351 InputValue::Length(2),
352 InputKind::Length(1, 10),
353 ),
354 StudyInput::new(
355 "color",
356 "Colour",
357 InputValue::Color(Color32::RED),
358 InputKind::Color,
359 ),
360 StudyInput::new(
361 "source",
362 "Source",
363 InputValue::Source(StudySource::Price(PriceSource::HL2)),
364 InputKind::Source,
365 ),
366 StudyInput::new(
367 "interval",
368 "Interval",
369 InputValue::Timeframe(Some(Timeframe::Hour1)),
370 InputKind::Timeframe,
371 ),
372 StudyInput::new(
373 "symbol",
374 "Symbol",
375 InputValue::Symbol("TEST".into()),
376 InputKind::Symbol,
377 ),
378 ];
379 let mut study = CustomIndicator::new("Inputs", |_| Vec::new())
380 .with_inputs(schema, |bars, inputs| {
381 let InputValue::Length(length) = inputs[0].value else {
382 unreachable!()
383 };
384 assert!(
385 bars.iter()
386 .all(|bar| bar.payload.as_deref() == Some(&serde_json::json!(42)))
387 );
388 bars.iter()
389 .map(|bar| IndicatorValue::Single(bar.close * length as f64))
390 .collect()
391 })
392 .unwrap();
393 let mut bar = Bar::new(chrono::DateTime::UNIX_EPOCH, 3.0, 3.0, 3.0, 3.0, 1.0);
394 bar.payload = Some(Arc::new(serde_json::json!(42)));
395 study.calculate(std::slice::from_ref(&bar));
396 for (id, bad) in [
397 ("length", InputValue::Length(0)),
398 ("length", InputValue::Number(2.0)),
399 ("symbol", InputValue::Symbol("\n".into())),
400 ] {
401 assert!(study.set_input(id, bad).is_err());
402 assert_eq!(study.values().value(0, 0), Some(6.0));
403 }
404 study.set_input("length", InputValue::Length(4)).unwrap();
405 assert!(study.values().is_empty());
406 assert_eq!(study.inputs()[0].default, InputValue::Length(2));
407 let saved: Vec<_> = study
408 .inputs()
409 .into_iter()
410 .map(|input| (input.id, input.value))
411 .collect();
412 let restored: Vec<(String, InputValue)> =
413 serde_json::from_str(&serde_json::to_string(&saved).unwrap()).unwrap();
414 let mut cloned = study.clone_box();
415 for (id, value) in restored {
416 cloned.set_input(&id, value).unwrap();
417 }
418 cloned.calculate(std::slice::from_ref(&bar));
419 assert_eq!(cloned.values().value(0, 0), Some(12.0));
420 cloned.update_source(&[bar], 0, &[5.0]);
421 assert_eq!(cloned.values().value(0, 0), Some(20.0));
422 let length = |value| StudyInput::new("length", "Length", value, InputKind::Length(1, 10));
423 for (schema, expected) in [
424 (
425 vec![StudyInput::new(
426 "",
427 "Length",
428 InputValue::Length(2),
429 InputKind::Length(1, 10),
430 )],
431 "study input 0 has an empty id",
432 ),
433 (
434 vec![length(InputValue::Length(2)), length(InputValue::Length(3))],
435 "study input \"length\" is declared twice",
436 ),
437 (
438 vec![length(InputValue::Length(0))],
439 "study input \"length\" has value Length(0), expected Length(1, 10)",
440 ),
441 (
442 vec![length(InputValue::Length(2)).with_default(InputValue::Length(11))],
443 "study input \"length\" has default Length(11), expected Length(1, 10)",
444 ),
445 ] {
446 let error = CustomIndicator::new("Inputs", |_| Vec::new())
447 .with_inputs(schema, |_, _| Vec::new())
448 .err()
449 .unwrap();
450 assert_eq!(error.to_string(), expected);
451 }
452 }
453
454 #[test]
455 fn clone_keeps_the_calculation() {
456 let indicator = CustomIndicator::new("Constant", |_| vec![IndicatorValue::Single(1.0); 3]);
457
458 let mut cloned = indicator.clone_box();
459 cloned.calculate(&[]);
460 assert_eq!(cloned.values().len(), 3);
461 }
462}