ArgByteOrder

Author: youknowone

vm/src/builtins/int.rs

@@ -5,7 +5,10 @@ use crate::{
     common::hash,
     convert::{ToPyObject, ToPyResult},
     format::FormatSpec,
-    function::{ArgIntoBool, OptionalArg, OptionalOption, PyArithmeticValue, PyComparisonValue},
+    function::{
+        ArgByteOrder, ArgIntoBool, OptionalArg, OptionalOption, PyArithmeticValue,
+        PyComparisonValue,
+    },
     types::{Comparable, Constructor, Hashable, PyComparisonOp},
     AsObject, Context, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromBorrowedObject,
     VirtualMachine,
@@ -626,15 +629,12 @@ impl PyInt {
         vm: &VirtualMachine,
     ) -> PyResult<PyRef<Self>> {
         let signed = args.signed.map_or(false, Into::into);
-        let value = match (args.byteorder.as_str(), signed) {
-            ("big", true) => BigInt::from_signed_bytes_be(&args.bytes.elements),
-            ("big", false) => BigInt::from_bytes_be(Sign::Plus, &args.bytes.elements),
-            ("little", true) => BigInt::from_signed_bytes_le(&args.bytes.elements),
-            ("little", false) => BigInt::from_bytes_le(Sign::Plus, &args.bytes.elements),
-            _ => {
-                return Err(
-                    vm.new_value_error("byteorder must be either 'little' or 'big'".to_owned())
-                )
+        let value = match (args.byteorder, signed) {
+            (ArgByteOrder::Big, true) => BigInt::from_signed_bytes_be(&args.bytes.elements),
+            (ArgByteOrder::Big, false) => BigInt::from_bytes_be(Sign::Plus, &args.bytes.elements),
+            (ArgByteOrder::Little, true) => BigInt::from_signed_bytes_le(&args.bytes.elements),
+            (ArgByteOrder::Little, false) => {
+                BigInt::from_bytes_le(Sign::Plus, &args.bytes.elements)
             }
         };
         Self::with_value(cls, value, vm)
@@ -656,16 +656,11 @@ impl PyInt {
             _ => {}
         }
 
-        let mut origin_bytes = match (args.byteorder.as_str(), signed) {
-            ("big", true) => value.to_signed_bytes_be(),
-            ("big", false) => value.to_bytes_be().1,
-            ("little", true) => value.to_signed_bytes_le(),
-            ("little", false) => value.to_bytes_le().1,
-            _ => {
-                return Err(
-                    vm.new_value_error("byteorder must be either 'little' or 'big'".to_owned())
-                );
-            }
+        let mut origin_bytes = match (args.byteorder, signed) {
+            (ArgByteOrder::Big, true) => value.to_signed_bytes_be(),
+            (ArgByteOrder::Big, false) => value.to_bytes_be().1,
+            (ArgByteOrder::Little, true) => value.to_signed_bytes_le(),
+            (ArgByteOrder::Little, false) => value.to_bytes_le().1,
         };
 
         let origin_len = origin_bytes.len();
@@ -678,21 +673,21 @@ impl PyInt {
             _ => vec![0u8; byte_len - origin_len],
         };
 
-        let bytes = match args.byteorder.as_str() {
-            "big" => {
+        let bytes = match args.byteorder {
+            ArgByteOrder::Big => {
                 let mut bytes = append_bytes;
                 bytes.append(&mut origin_bytes);
                 bytes
             }
-            "little" => {
+            ArgByteOrder::Little => {
                 let mut bytes = origin_bytes;
                 bytes.append(&mut append_bytes);
                 bytes
             }
-            _ => Vec::new(),
         };
         Ok(bytes.into())
     }
+
     #[pyproperty]
     fn real(&self, vm: &VirtualMachine) -> PyRef<Self> {
         // subclasses must return int here
@@ -759,15 +754,15 @@ pub struct IntOptions {
 #[derive(FromArgs)]
 struct IntFromByteArgs {
     bytes: PyBytesInner,
-    byteorder: PyStrRef,
+    byteorder: ArgByteOrder,
     #[pyarg(named, optional)]
     signed: OptionalArg<ArgIntoBool>,
 }
 
 #[derive(FromArgs)]
 struct IntToByteArgs {
     length: PyIntRef,
-    byteorder: PyStrRef,
+    byteorder: ArgByteOrder,
     #[pyarg(named, optional)]
     signed: OptionalArg<ArgIntoBool>,
 }

vm/src/function/mod.rs

@@ -16,3 +16,26 @@ pub use builtin::{IntoPyNativeFunc, OwnedParam, PyNativeFunc, RefParam};
 pub use either::Either;
 pub use number::{ArgIntoBool, ArgIntoComplex, ArgIntoFloat};
 pub use protocol::{ArgCallable, ArgIterable, ArgMapping, ArgSequence};
+
+use crate::{builtins::PyStr, convert::TryFromBorrowedObject, PyObject, PyResult, VirtualMachine};
+
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub enum ArgByteOrder {
+    Big,
+    Little,
+}
+
+impl TryFromBorrowedObject for ArgByteOrder {
+    fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObject) -> PyResult<Self> {
+        obj.try_value_with(
+            |s: &PyStr| match s.as_str() {
+                "big" => Ok(Self::Big),
+                "little" => Ok(Self::Little),
+                _ => {
+                    Err(vm.new_value_error("byteorder must be either 'little' or 'big'".to_owned()))
+                }
+            },
+            vm,
+        )
+    }
+}