Skip to content

Commit

Permalink
refactor: move more metadata handling into meta
Browse files Browse the repository at this point in the history
  • Loading branch information
agoose77 committed Nov 20, 2023
1 parent 4873236 commit d0a0026
Show file tree
Hide file tree
Showing 4 changed files with 59 additions and 14 deletions.
12 changes: 6 additions & 6 deletions src/awkward/_meta/recordmeta.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,11 +71,11 @@ def dimension_optiontype(self) -> bool:
return False

@property
def is_leaf(self):
def is_leaf(self) -> bool:
return len(self._contents) == 0

@property
def contents(self):
def contents(self) -> list[T]:
return self._contents

@property
Expand All @@ -85,7 +85,7 @@ def fields(self) -> list[str]:
else:
return self._fields

def index_to_field(self, index):
def index_to_field(self, index: int) -> str:
if 0 <= index < len(self._contents):
if self._fields is None:
return str(index)
Expand All @@ -96,7 +96,7 @@ def index_to_field(self, index):
f"no index {index} in record with {len(self._contents)} fields"
)

def field_to_index(self, field):
def field_to_index(self, field: str) -> int:
if self._fields is None:
try:
i = int(field)
Expand All @@ -116,7 +116,7 @@ def field_to_index(self, field):
f"no field {field!r} in record with {len(self._contents)} fields"
)

def has_field(self, field):
def has_field(self, field: str) -> bool:
if self._fields is None:
try:
i = int(field)
Expand All @@ -127,7 +127,7 @@ def has_field(self, field):
else:
return field in self._fields

def content(self, index_or_field):
def content(self, index_or_field: int | str) -> T:
if is_integer(index_or_field):
index = index_or_field
elif isinstance(index_or_field, str):
Expand Down
3 changes: 3 additions & 0 deletions src/awkward/_meta/unionmeta.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,3 +94,6 @@ def dimension_optiontype(self) -> bool:
if content.dimension_optiontype:
return True
return False

def content(self, index: int) -> T:
return self._contents[index]
3 changes: 0 additions & 3 deletions src/awkward/forms/unionform.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,9 +108,6 @@ def _union_of_optionarrays(self, index, parameters):
.form
)

def content(self, index):
return self._contents[index]

def __repr__(self):
args = [
repr(self._tags),
Expand Down
55 changes: 50 additions & 5 deletions src/awkward/types/recordtype.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,10 @@
import awkward._prettyprint
from awkward._behavior import find_record_typestr
from awkward._parameters import parameters_are_equal, type_parameters_equal
from awkward._regularize import is_integer
from awkward._typing import Any, JSONMapping, Self, cast, final
from awkward._util import UNSET, Sentinel
from awkward.errors import FieldNotFoundError
from awkward.types.type import Type


Expand Down Expand Up @@ -225,13 +227,56 @@ def _is_equal_to(self, other: Any, all_parameters: bool) -> bool:
return False

def index_to_field(self, index: int) -> str:
return ak.forms.RecordForm.index_to_field(self, index) # type: ignore[arg-type]
if 0 <= index < len(self._contents):
if self._fields is None:
return str(index)
else:
return self._fields[index]
else:
raise IndexError(
f"no index {index} in record with {len(self._contents)} fields"
)

def field_to_index(self, field: str) -> int:
return ak.forms.RecordForm.field_to_index(self, field) # type: ignore[arg-type]
if self._fields is None:
try:
i = int(field)
except ValueError:
pass
else:
if 0 <= i < len(self._contents):
return i
else:
try:
i = self._fields.index(field)
except ValueError:
pass
else:
return i
raise FieldNotFoundError(
f"no field {field!r} in record with {len(self._contents)} fields"
)

def has_field(self, field: str) -> bool:
return ak.forms.RecordForm.has_field(self, field) # type: ignore[arg-type]
if self._fields is None:
try:
i = int(field)
except ValueError:
return False
else:
return 0 <= i < len(self._contents)
else:
return field in self._fields

def content(self, index_or_field: str | int) -> Type:
return ak.forms.RecordForm.content(self, index_or_field) # type: ignore[arg-type]
def content(self, index_or_field: int | str) -> Type:
if is_integer(index_or_field):
index = int(index_or_field)
elif isinstance(index_or_field, str):
index = self.field_to_index(index_or_field)
else:
raise TypeError(
"index_or_field must be an integer (index) or string (field), not {}".format(
repr(index_or_field)
)
)
return self._contents[index]

0 comments on commit d0a0026

Please sign in to comment.