From f820dfcffa33bb6a7988434b8454538d66b85aff Mon Sep 17 00:00:00 2001
From: Jaco Greeff <jacogr@gmail.com>
Date: Tue, 19 Feb 2019 15:20:45 +0100
Subject: [PATCH] Inject index properly into ReferendumInfo (with unwrap)
 (#718)

* Inject inject properly into ReferendumInfo (with unwrap)

* Use unwrapOr

* Fix type import
---
 .../src/democracy/referendumInfo.ts           | 23 ++++++++++---------
 packages/types/src/type/ReferendumInfo.ts     |  6 +++--
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/packages/api-derive/src/democracy/referendumInfo.ts b/packages/api-derive/src/democracy/referendumInfo.ts
index 243d461e9e..4897d7a8f8 100644
--- a/packages/api-derive/src/democracy/referendumInfo.ts
+++ b/packages/api-derive/src/democracy/referendumInfo.ts
@@ -16,10 +16,10 @@ import { drr } from '../util/drr';
  * A [[ReferendumInfo]] with an additional `index` field
  */
 export class ReferendumInfoExtended extends ReferendumInfo {
-  constructor (info: any) {
-    super(info);
+  constructor (value: ReferendumInfo) {
+    super(value);
 
-    this.set('index', new ReferendumIndex(info.index));
+    this.set('index', value.get('index') as ReferendumIndex);
   }
 
   /**
@@ -34,14 +34,15 @@ export function referendumInfo (api: ApiInterface$Rx) {
   return (index: BN | number): Observable<Option<ReferendumInfoExtended>> => {
     return (api.query.democracy.referendumInfoOf(index) as Observable<Option<ReferendumInfo>>)
       .pipe(
-        map((info) =>
-          new Option<ReferendumInfoExtended>(
-            ReferendumInfoExtended,
-            info.isSome
-              ? { ...info.toJSON(), index }
-              : null
-          )
-        ),
+        map((optionInfo) => {
+          const info = optionInfo.unwrapOr(null);
+
+          if (info) {
+            info.set('index', new ReferendumIndex(index));
+          }
+
+          return new Option<ReferendumInfoExtended>(ReferendumInfoExtended, info);
+        }),
         drr()
       );
   };
diff --git a/packages/types/src/type/ReferendumInfo.ts b/packages/types/src/type/ReferendumInfo.ts
index c322b83271..efffb6d7b1 100644
--- a/packages/types/src/type/ReferendumInfo.ts
+++ b/packages/types/src/type/ReferendumInfo.ts
@@ -2,8 +2,10 @@
 // This software may be modified and distributed under the terms
 // of the Apache-2.0 license. See the LICENSE file for details.
 
-import BlockNumber from './BlockNumber';
+import { Codec } from '../types';
+
 import Struct from '../codec/Struct';
+import BlockNumber from './BlockNumber';
 import Proposal from './Proposal';
 import VoteThreshold from './VoteThreshold';
 
@@ -20,7 +22,7 @@ type ReferendumInfoValue = {
  * Info regarding an ongoing referendum
  */
 export default class ReferendumInfo extends Struct {
-  constructor (value?: ReferendumInfoValue | Uint8Array) {
+  constructor (value?: ReferendumInfoValue | Uint8Array | Map<string, Codec>) {
     super({
       end: BlockNumber,
       proposal: Proposal,