@@ -101,6 +101,100 @@ def add_dimension(self, name: str, size: int) -> "MDIODatasetBuilder":
101101 self ._state = _BuilderState .HAS_DIMENSIONS
102102 return self
103103
104+ def push_dimension (self , dimension : NamedDimension , position : int , new_dim_chunk_size : int = 1 ) -> "MDIODatasetBuilder" :
105+ """Pushes a dimension to all Coordiantes and Variables.
106+ The position argument is the domain index of the dimension to push.
107+ If a Variable is within the position domain, it will be inserted at the position and all remaining dimensions will be shifted to the right.
108+
109+ Args:
110+ dimension: The dimension to push
111+ position: The position to push the dimension to
112+ new_dim_chunk_size: The chunk size for only the new dimension
113+
114+ Returns:
115+ self: Returns self for method chaining
116+ """
117+ if position < 0 :
118+ msg = "Support for negative positions is not implemented yet!"
119+ raise ValueError (msg )
120+ if position > len (self ._dimensions ):
121+ msg = "Position is greater than the number of dimensions"
122+ raise ValueError (msg )
123+ if new_dim_chunk_size <= 0 :
124+ # TODO(BrianMichell): Do we actually need to check this, or does Pydantic handle when we call?
125+ msg = "New dimension chunk size must be greater than 0"
126+ raise ValueError (msg )
127+
128+ # print("###########################STATE BEFORE INSERTING DIMENSION ###########################")
129+ # for d in self._dimensions:
130+ # print(d.model_dump_json())
131+ # for c in self._coordinates:
132+ # print(c.model_dump_json())
133+ # for v in self._variables:
134+ # print(v.model_dump_json())
135+ # print("########################################################################################")
136+
137+
138+ # In-place insertion of the dimension to the existing list of dimensions
139+ self ._dimensions .insert (position , dimension )
140+
141+ # def propogate_dimension(variable: Variable, position: int, new_dim_chunk_size: int) -> Variable:
142+ # """Propogates the dimension to the variable or coordinate."""
143+ # if len(variable.dimensions) <= position:
144+ # # Don't do anything if the new dimension is not within the Variable's domain
145+ # return variable
146+ # if variable.name == "trace_mask":
147+ # # Special case for trace_mask. Don't do anything.
148+ # return variable
149+ # # new_dimensions = variable.dimensions[:position] + (dimension,) + variable.dimensions[position:]
150+ # # new_chunk_sizes = variable.chunk_sizes[:position] + (new_dim_chunk_size,) + variable.chunk_sizes[position:]
151+ # new_dimensions = variable.dimensions[:position] + [dimension] + variable.dimensions[position:]
152+ # new_chunk_sizes = variable.chunk_sizes[:position] + [new_dim_chunk_size] + variable.chunk_sizes[position:]
153+ # return variable.model_copy(update={"dimensions": new_dimensions, "chunk_sizes": new_chunk_sizes})
154+ def propogate_dimension (variable : Variable , position : int , new_dim_chunk_size : int ) -> Variable :
155+ """Propogates the dimension to the variable or coordinate."""
156+ from mdio .builder .schemas .chunk_grid import RegularChunkGrid , RegularChunkShape
157+ if len (variable .dimensions ) <= position :
158+ # Don't do anything if the new dimension is not within the Variable's domain
159+ return variable
160+ if variable .name == "trace_mask" :
161+ # Special case for trace_mask. Don't do anything.
162+ return variable
163+ # new_dimensions = variable.dimensions[:position] + (dimension,) + variable.dimensions[position:]
164+ # new_chunk_sizes = variable.chunk_sizes[:position] + (new_dim_chunk_size,) + variable.chunk_sizes[position:]
165+ new_dimensions = variable .dimensions [:position ] + [dimension ] + variable .dimensions [position :]
166+
167+ # Get current chunk shape from metadata
168+ current_chunk_shape = (1 ,) * len (variable .dimensions ) # Default fallback
169+ if variable .metadata is not None and variable .metadata .chunk_grid is not None :
170+ current_chunk_shape = variable .metadata .chunk_grid .configuration .chunk_shape
171+
172+ # Insert new chunk size at the correct position
173+ new_chunk_shape = current_chunk_shape [:position ] + (new_dim_chunk_size ,) + current_chunk_shape [position :]
174+
175+ # Create new chunk grid configuration
176+ new_chunk_grid = RegularChunkGrid (configuration = RegularChunkShape (chunk_shape = new_chunk_shape ))
177+
178+ # Update metadata with new chunk grid
179+ new_metadata = variable .metadata .model_copy () if variable .metadata else VariableMetadata ()
180+ new_metadata .chunk_grid = new_chunk_grid
181+ ret = variable .model_copy (update = {"dimensions" : new_dimensions , "metadata" : new_metadata })
182+ return ret
183+
184+ to_ignore = []
185+ for v in self ._dimensions :
186+ to_ignore .append (v .name )
187+ for c in self ._coordinates :
188+ to_ignore .append (c .name )
189+
190+ for i in range (len (self ._variables )):
191+ var = self ._variables [i ]
192+ if var .name in to_ignore :
193+ continue
194+ self ._variables [i ] = propogate_dimension (var , position , new_dim_chunk_size )
195+
196+ return self
197+
104198 def add_coordinate ( # noqa: PLR0913
105199 self ,
106200 name : str ,
0 commit comments