Coverage for src/cc_liquid/data_loader.py: 79%

1"""Data loading and abstraction for cc-liquid."""

2

3 from abc import ABC, abstractmethod

4

5 import polars as pl

6

7

8 class DataSource(ABC):

9 """Abstract base class for data sources."""

10

11 @abstractmethod

12 def load(self) -> pl.DataFrame:

13 """Load data into a Polars DataFrame."""

14 pass

15

16

17 class FileDataSource(DataSource):

18 """Loads prediction data from a Parquet or CSV file."""

19

20 def __init__(

21 self,

22 path: str,

23 date_column: str,

24 asset_id_column: str,

25 prediction_column: str,

26 ):

27 self.path = path

28 self.date_column = date_column

29 self.asset_id_column = asset_id_column

30 self.prediction_column = prediction_column

31

32 def load(self) -> pl.DataFrame:

33 """Loads data from the file."""

34 if self.path.endswith(".parquet"):

35 df = pl.read_parquet(self.path)

36 elif self.path.endswith(".csv"):

37 df = pl.read_csv(self.path)

38 else:

39 raise ValueError("Unsupported file type. Use .parquet or .csv.")

40

41 return df

42

43

44 class DataFrameDataSource(DataSource):

45 """Uses an existing Polars DataFrame as the data source."""

46

47 def __init__(

48 self,

49 df: pl.DataFrame,

50 date_column: str,

51 asset_id_column: str,

52 prediction_column: str,

53 ):

54 self.df = df

55 self.date_column = date_column

56 self.asset_id_column = asset_id_column

57 self.prediction_column = prediction_column

58

59 def load(self) -> pl.DataFrame:

60 """Returns the existing DataFrame."""

61 return self.df

62

63

64 class DataLoader:

65 """Factory for creating data sources."""

66

67 @staticmethod

68 def from_file(path: str, date_col: str, id_col: str, pred_col: str) -> pl.DataFrame:

69 """Create a file data source and load data."""

70 return FileDataSource(

71 path,

72 date_column=date_col,

73 asset_id_column=id_col,

74 prediction_column=pred_col,

75 ).load()

76

77 @staticmethod

78 def from_dataframe(

79 df: pl.DataFrame, date_col: str, id_col: str, pred_col: str

80 ) -> pl.DataFrame:

81 """Create a DataFrame data source and load data."""

82 return DataFrameDataSource(

83 df,

84 date_column=date_col,

85 asset_id_column=id_col,

86 prediction_column=pred_col,

87 ).load()

88

89 @staticmethod

90 def from_crowdcent_api(

91 api_key: str | None = None,

92 challenge_slug: str = "hyperliquid-ranking",

93 download_path: str | None = None,

94 date_col: str = "release_date",

95 id_col: str = "id",

96 pred_col: str = "pred_10d",

97 ) -> pl.DataFrame:

98 """

99 Download and load the CrowdCent meta model.

100

101 Args:

102 api_key: CrowdCent API key (if None, will try to load from env)

103 challenge_slug: The challenge to download data for

104 download_path: Optional path to save the downloaded file

105 date_col: Date column name in the meta model

106 id_col: Asset ID column name in the meta model

107 pred_col: Prediction column name to use from the meta model

108

109 Returns:

110 Polars DataFrame with original column names

111 """

112 from crowdcent_challenge import ChallengeClient

113

114 if api_key is None:

115 import os

116

117 api_key = os.getenv("CROWDCENT_API_KEY")

118 if not api_key:

119 raise ValueError("CROWDCENT_API_KEY not found in environment variables")

120

121 client = ChallengeClient(challenge_slug=challenge_slug, api_key=api_key)

122

123 if download_path is None:

124 download_path = "predictions.parquet"

125

126 client.download_meta_model(download_path)

127

128 return DataLoader.from_file(

129 path=download_path, date_col=date_col, id_col=id_col, pred_col=pred_col

130 )

131

132 @staticmethod

133 def from_numerai_api(

134 download_path: str | None = None,

135 date_col: str = "date",

136 id_col: str = "symbol",

137 pred_col: str = "meta_model",

138 ) -> pl.DataFrame:

139 """

140 Download and load the Numerai crypto meta model.

141

142 Args:

143 download_path: Optional path to save the downloaded file

144 date_col: Date column name in the meta model

145 id_col: Asset ID/symbol column name in the meta model

146 pred_col: Prediction column name to use from the meta model

147

148 Returns:

149 Polars DataFrame with original column names

150 """

151 try:

152 from numerapi import CryptoAPI

153 except ImportError:

154 raise ImportError(

155 "numerapi is required. Install with: uv add cc-liquid[numerai]"

156 )

157

158 api = CryptoAPI()

159

160 if download_path is None:

161 download_path = "predictions.parquet"

162

163 api.download_dataset("v1.0/historical_meta_models.parquet", download_path)

164

165 return DataLoader.from_file(

166 path=download_path, date_col=date_col, id_col=id_col, pred_col=pred_col

167 )