TBaseNumber to std::ostream print support.

2 files changed, 59 insertions, 4 deletions

diff --git a/util/stream/format.h b/util/stream/format.h
index b033208a1b..daba3b8a05 100644
--- a/util/stream/format.h
+++ b/util/stream/format.h
@@ -28,12 +28,28 @@ namespace NFormatPrivate {
     template <>
     struct TLog2<1>: std::integral_constant<size_t, 0> {};
 
-    static inline void WriteChars(IOutputStream& os, char c, size_t count) {
+    template <typename T>
+    inline void StreamWrite(T& stream, const char* s, size_t size) {
+        stream.write(s, size);
+    }
+
+    template <>
+    inline void StreamWrite(IOutputStream& stream, const char* s, size_t size) {
+        stream.Write(s, size);
+    }
+
+    template <>
+    inline void StreamWrite(TStringStream& stream, const char* s, size_t size) {
+        stream.Write(s, size);
+    }
+
+    template <typename T>
+    static inline void WriteChars(T& os, char c, size_t count) {
         if (count == 0)
             return;
         TTempBuf buf(count);
         memset(buf.Data(), c, count);
-        os.Write(buf.Data(), count);
+        StreamWrite(os, buf.Data(), count);
     }
 
     template <typename T>
@@ -106,8 +122,8 @@ namespace NFormatPrivate {
     template <typename T, size_t Base>
     using TUnsignedBaseNumber = TBaseNumber<std::make_unsigned_t<std::remove_cv_t<T>>, Base>;
 
-    template <typename T, size_t Base>
-    IOutputStream& operator<<(IOutputStream& stream, const TBaseNumber<T, Base>& value) {
+    template <typename TStream, typename T, size_t Base>
+    TStream& ToStreamImpl(TStream& stream, const TBaseNumber<T, Base>& value) {
         char buf[8 * sizeof(T) + 1]; /* Add 1 for sign. */
         TStringBuf str(buf, IntToString<Base>(value.Value, buf, sizeof(buf)));
 
@@ -132,6 +148,16 @@ namespace NFormatPrivate {
         return stream;
     }
 
+    template <typename T, size_t Base>
+    IOutputStream& operator<<(IOutputStream& stream, const TBaseNumber<T, Base>& value) {
+        return ToStreamImpl(stream, value);
+    }
+
+    template <typename T, size_t Base>
+    std::ostream& operator<<(std::ostream& stream, const TBaseNumber<T, Base>& value) {
+        return ToStreamImpl(stream, value);
+    }
+
     template <typename Char, size_t Base>
     struct TBaseText {
         TBasicStringBuf<Char> Text;
diff --git a/util/stream/format_std_ut.cpp b/util/stream/format_std_ut.cpp
new file mode 100644
index 0000000000..b903000382
--- /dev/null
+++ b/util/stream/format_std_ut.cpp
@@ -0,0 +1,29 @@
+#include "format.h"
+
+#include <library/cpp/testing/unittest/registar.h>
+
+#include <sstream>
+
+Y_UNIT_TEST_SUITE(StdOstreamFormattingTest) {
+    template <typename T>
+    TString ToStringViaOstream(T baseNumber) {
+        std::stringstream ss;
+        ss << baseNumber;
+        return ss.str();
+    }
+
+    Y_UNIT_TEST(TestBin) {
+        UNIT_ASSERT_VALUES_EQUAL(ToStringViaOstream(Bin(static_cast<ui32>(2), nullptr)), "10");
+        UNIT_ASSERT_VALUES_EQUAL(ToStringViaOstream(SBin(static_cast<i32>(-2), nullptr)), "-10");
+        UNIT_ASSERT_VALUES_EQUAL(ToStringViaOstream(SBin(static_cast<i32>(-2))), "-0b00000000000000000000000000000010");
+        UNIT_ASSERT_VALUES_EQUAL(ToStringViaOstream(SBin(static_cast<i32>(-2), HF_FULL)), "-00000000000000000000000000000010");
+        UNIT_ASSERT_VALUES_EQUAL(ToStringViaOstream(Bin(static_cast<ui32>(15), nullptr)), "1111");
+        UNIT_ASSERT_VALUES_EQUAL(ToStringViaOstream(Bin(static_cast<ui32>(1))), "0b00000000000000000000000000000001");
+        UNIT_ASSERT_VALUES_EQUAL(ToStringViaOstream(Bin(static_cast<ui32>(-1))), "0b11111111111111111111111111111111");
+        UNIT_ASSERT_VALUES_EQUAL(ToStringViaOstream(Bin(static_cast<i32>(-1))), "0b11111111111111111111111111111111");
+        UNIT_ASSERT_VALUES_EQUAL(ToStringViaOstream(Bin(static_cast<i32>(-1), nullptr)), "11111111111111111111111111111111");
+        UNIT_ASSERT_VALUES_EQUAL(ToStringViaOstream(Bin(static_cast<ui32>(256))), "0b00000000000000000000000100000000");
+        UNIT_ASSERT_VALUES_EQUAL(ToStringViaOstream(Bin(static_cast<ui8>(16))), "0b00010000");
+        UNIT_ASSERT_VALUES_EQUAL(ToStringViaOstream(Bin(static_cast<ui64>(1234587912357ull))), "0b0000000000000000000000010001111101110011001011001000100010100101");
+    }
+}