Coverage for src/pystencilssfg/lang/cpp/std_span.py: 85%

34 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-04-04 07:16 +0000

1from pystencils import Field, DynamicType 

2from pystencils.types import UserTypeSpec, create_type, PsType 

3 

4from ...lang import SupportsFieldExtraction, AugExpr, cpptype 

5 

6 

7class StdSpan(AugExpr, SupportsFieldExtraction): 

8 _template = cpptype("std::span< {T} >", "<span>") 

9 

10 def __init__(self, T: UserTypeSpec, ref=False, const=False): 

11 T = create_type(T) 

12 dtype = self._template(T=T, const=const, ref=ref) 

13 self._element_type = T 

14 super().__init__(dtype) 

15 

16 @property 

17 def element_type(self) -> PsType: 

18 return self._element_type 

19 

20 def _extract_ptr(self) -> AugExpr: 

21 return AugExpr.format("{}.data()", self) 

22 

23 def _extract_size(self, coordinate: int) -> AugExpr | None: 

24 if coordinate > 0: 

25 return None 

26 else: 

27 return AugExpr.format("{}.size()", self) 

28 

29 def _extract_stride(self, coordinate: int) -> AugExpr | None: 

30 if coordinate > 0: 

31 return None 

32 else: 

33 return AugExpr.format("1") 

34 

35 @staticmethod 

36 def from_field(field: Field, ref: bool = False, const: bool = False): 

37 if field.spatial_dimensions > 1 or field.index_shape not in ((), (1,)): 

38 raise ValueError( 

39 "Only one-dimensional fields with trivial index dimensions can be mapped onto `std::span`" 

40 ) 

41 if isinstance(field.dtype, DynamicType): 

42 raise ValueError("Cannot map dynamically typed field to std::span") 

43 

44 return StdSpan(field.dtype, ref=ref, const=const).var(field.name) 

45 

46 

47def std_span_ref(field: Field): 

48 from warnings import warn 

49 

50 warn( 

51 "`std_span_ref` is deprecated and will be removed in version 0.1. Use `std.span.from_field` instead.", 

52 FutureWarning, 

53 ) 

54 return StdSpan.from_field(field, ref=True)